Method forward

models/hidream.py:364–411

(self, inputs)

    @torch.autocast('cuda', dtype=AUTOCAST_DTYPE)
    def forward(self, inputs):
        hidden_states, img_ids, timesteps, pooled_embeds, t5_prompt_embeds, llama3_prompt_embeds = inputs

        batch_size = hidden_states.shape[0]
        hidden_states_type = hidden_states.dtype

        timesteps = self.expand_timesteps(timesteps, batch_size, hidden_states.device)
        timesteps = self.t_embedder(timesteps, hidden_states_type)
        p_embedder = self.p_embedder(pooled_embeds)
        adaln_input = timesteps + p_embedder

        hidden_states = self.x_embedder(hidden_states)

        T5_encoder_hidden_states = t5_prompt_embeds
        encoder_hidden_states = llama3_prompt_embeds
        encoder_hidden_states = [encoder_hidden_states[k] for k in self.llama_layers]

        assert self.caption_projection is not None
        if self.caption_projection is not None:
            new_encoder_hidden_states = []
            for i, enc_hidden_state in enumerate(encoder_hidden_states):
                enc_hidden_state = self.caption_projection[i](enc_hidden_state)
                enc_hidden_state = enc_hidden_state.view(batch_size, -1, hidden_states.shape[-1])
                new_encoder_hidden_states.append(enc_hidden_state)
            encoder_hidden_states = new_encoder_hidden_states
            T5_encoder_hidden_states = self.caption_projection[-1](T5_encoder_hidden_states)
            T5_encoder_hidden_states = T5_encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1])
            encoder_hidden_states.append(T5_encoder_hidden_states)

        txt_ids = torch.zeros(
            batch_size,
            encoder_hidden_states[-1].shape[1] + encoder_hidden_states[-2].shape[1] + encoder_hidden_states[0].shape[1],
            3,
            device=img_ids.device, dtype=img_ids.dtype
        )
        ids = torch.cat((img_ids, txt_ids), dim=1)
        rope = self.pe_embedder(ids)

        initial_encoder_hidden_states = torch.cat([encoder_hidden_states[-1], encoder_hidden_states[-2]], dim=1)
        llama_encoder_hidden_states = torch.stack(encoder_hidden_states[:-1], dim=0)

        # With nf4 quantization, tensors can end up float32, which breaks flash attention later, so we cast it here.
        hidden_states = hidden_states.to(AUTOCAST_DTYPE)
        initial_encoder_hidden_states = initial_encoder_hidden_states.to(AUTOCAST_DTYPE)
        llama_encoder_hidden_states = llama_encoder_hidden_states.to(AUTOCAST_DTYPE)
        adaln_input = adaln_input.to(AUTOCAST_DTYPE)

        return make_contiguous(hidden_states, initial_encoder_hidden_states, llama_encoder_hidden_states, adaln_input, rope)


class TransformerWrapper(nn.Module):