FLUX - Was Dumbed Down - We Can Restore It

This has been a topic of discussion in post and DM with AbstractPilia and we have been working different approaches to fix this issue.

My focus has been on the timestep embedding and why a deterministic script was used:

If we call the FLUX timestep embedding rather then a sinusoidal timestep embedding, the answer seems clear.

The Timestep embed needs to be in Torch Long INT64 (I wondered why comfy added 64bit)

If calling the FLUX timestep embedding the model needs to be on one device, this would limit the number of users severally and may not even fit into a 24GB card.

For High VRAM users try:

# Replace the sinusoidal embedding with learnable embeddings
class TimeEmbedding(nn.Module):
    def __init__(self, max_timestep: int, embedding_dim: int):
        """
        Create learnable time embeddings.
        :param max_timestep: The maximum number of timesteps.
        :param embedding_dim: The dimension of the output embeddings.
        """
        super().__init__()
        # Define the learnable embedding layer for time steps
        self.embedding = nn.Embedding(max_timestep, embedding_dim)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """
        Get time embeddings for the provided time indices.
        :param t: A tensor of time indices.
        :return: The corresponding time embeddings.
        """
        # Ensure that the tensor `t` is cast to Long (torch.int64) before passing to the embedding layer
        t = t.to(torch.long)
        return self.embedding(t) 

It is possible that a different approach could be used for this, but this is beyond my level.