Class UNet2DConditionModelXL

python_coreml_stable_diffusion/unet.py:1041–1142

UNet2DConditionModel variant for Stable Diffusion XL with an extended forward() signature

class UNet2DConditionModelXL(UNet2DConditionModel):
    """ UNet2DConditionModel variant for Stable Diffusion XL
    with an extended forward() signature
    """
    def forward(
        self,
        sample,
        timestep,
        encoder_hidden_states,
        time_ids,
        text_embeds,
        *additional_residuals,
    ):
        # 0. Project time embeddings
        t_emb = self.time_proj(timestep)
        emb = self.time_embedding(t_emb)

        aug_emb = None

        if self.config.addition_embed_type == "text":
            raise NotImplementedError
        elif self.config.addition_embed_type == "text_image":
            raise NotImplementedError
        elif self.config.addition_embed_type == "text_time":
            assert time_ids is not None
            assert text_embeds is not None

            time_embeds = self.add_time_proj(time_ids.flatten())
            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))

            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
            aug_emb = self.add_embedding(add_embeds)
        elif self.config.addition_embed_type == "image":
            raise NotImplementedError
        elif self.config.addition_embed_type == "image_hint":
            raise NotImplementedError

        emb = emb + aug_emb if aug_emb is not None else emb

        # 1. center input if necessary
        if self.config.center_input_sample:
            sample = 2 * sample - 1.0

        # 2. pre-process
        sample = self.conv_in(sample)

        # 3. down
        down_block_res_samples = (sample, )
        for downsample_block in self.down_blocks:
            if hasattr(
                    downsample_block,
                    "attentions") and downsample_block.attentions is not None:
                sample, res_samples = downsample_block(
                    hidden_states=sample,
                    temb=emb,
                    encoder_hidden_states=encoder_hidden_states)
            else:
                sample, res_samples = downsample_block(hidden_states=sample,