Method forward

diffsynth/models/sdxl_unet.py:88–136

(
        self,
        sample, timestep, encoder_hidden_states, add_time_id, add_text_embeds,
        tiled=False, tile_size=64, tile_stride=8,
        use_gradient_checkpointing=False,
        **kwargs
    )

        self.is_kolors = is_kolors

    def forward(
        self,
        sample, timestep, encoder_hidden_states, add_time_id, add_text_embeds,
        tiled=False, tile_size=64, tile_stride=8,
        use_gradient_checkpointing=False,
        **kwargs
    ):
        # 1. time
        t_emb = self.time_proj(timestep).to(sample.dtype)
        t_emb = self.time_embedding(t_emb)
        
        time_embeds = self.add_time_proj(add_time_id)
        time_embeds = time_embeds.reshape((add_text_embeds.shape[0], -1))
        add_embeds = torch.concat([add_text_embeds, time_embeds], dim=-1)
        add_embeds = add_embeds.to(sample.dtype)
        add_embeds = self.add_time_embedding(add_embeds)

        time_emb = t_emb + add_embeds
        
        # 2. pre-process
        height, width = sample.shape[2], sample.shape[3]
        hidden_states = self.conv_in(sample)
        text_emb = encoder_hidden_states if self.text_intermediate_proj is None else self.text_intermediate_proj(encoder_hidden_states)
        res_stack = [hidden_states]
        
        # 3. blocks
        def create_custom_forward(module):
            def custom_forward(*inputs):
                return module(*inputs)
            return custom_forward
        for i, block in enumerate(self.blocks):
            if self.training and use_gradient_checkpointing and not (isinstance(block, PushBlock) or isinstance(block, PopBlock)):
                hidden_states, time_emb, text_emb, res_stack = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states, time_emb, text_emb, res_stack,
                    use_reentrant=False,
                )
            else:
                hidden_states, time_emb, text_emb, res_stack = block(
                    hidden_states, time_emb, text_emb, res_stack,
                    tiled=tiled, tile_size=tile_size, tile_stride=tile_stride
                )
        
        # 4. output
        hidden_states = self.conv_norm_out(hidden_states)
        hidden_states = self.conv_act(hidden_states)
        hidden_states = self.conv_out(hidden_states)

        return hidden_states
    
    @staticmethod
    def state_dict_converter():