toolkit/models/flux.py

# forward that bypasses the guidance embedding so it can be avoided during training.
from functools import partial
from typing import Optional
import torch
from diffusers import FluxTransformer2DModel


def guidance_embed_bypass_forward(self, timestep, guidance, pooled_projection):
    timesteps_proj = self.time_proj(timestep)
    timesteps_emb = self.timestep_embedder(
        timesteps_proj.to(dtype=pooled_projection.dtype))  # (N, D)
    pooled_projections = self.text_embedder(pooled_projection)
    conditioning = timesteps_emb + pooled_projections
    return conditioning

# bypass the forward function


def bypass_flux_guidance(transformer):
    if hasattr(transformer.time_text_embed, '_bfg_orig_forward'):
        return
    # dont bypass if it doesnt have the guidance embedding
    if not hasattr(transformer.time_text_embed, 'guidance_embedder'):
        return
    transformer.time_text_embed._bfg_orig_forward = transformer.time_text_embed.forward
    transformer.time_text_embed.forward = partial(
        guidance_embed_bypass_forward, transformer.time_text_embed
    )

# restore the forward function


def restore_flux_guidance(transformer):
    if not hasattr(transformer.time_text_embed, '_bfg_orig_forward'):
        return
    transformer.time_text_embed.forward = transformer.time_text_embed._bfg_orig_forward
    del transformer.time_text_embed._bfg_orig_forward

def new_device_to(self: FluxTransformer2DModel, *args, **kwargs):
    # Store original device if provided in args or kwargs
    device_in_kwargs = 'device' in kwargs
    device_in_args = any(isinstance(arg, (str, torch.device)) for arg in args)
    
    device = None
    # Remove device from kwargs if present
    if device_in_kwargs:
        device = kwargs['device']
        del kwargs['device']
    
    # Only filter args if we detected a device argument
    if device_in_args:
        args = list(args)
        for idx, arg in enumerate(args):
            if isinstance(arg, (str, torch.device)):
                device = arg
                del args[idx]
    
    self.pos_embed = self.pos_embed.to(device, *args, **kwargs)
    self.time_text_embed = self.time_text_embed.to(device, *args, **kwargs)
    self.context_embedder = self.context_embedder.to(device, *args, **kwargs)
    self.x_embedder = self.x_embedder.to(device, *args, **kwargs)
    for block in self.transformer_blocks:
        block.to(block._split_device, *args, **kwargs)
    for block in self.single_transformer_blocks:
        block.to(block._split_device, *args, **kwargs)
    
    self.norm_out = self.norm_out.to(device, *args, **kwargs)
    self.proj_out = self.proj_out.to(device, *args, **kwargs)
    
    
    
    return self

    


def split_gpu_double_block_forward(
    self,
    hidden_states: torch.FloatTensor,
    encoder_hidden_states: torch.FloatTensor,
    temb: torch.FloatTensor,
    image_rotary_emb=None,
    joint_attention_kwargs=None,
):
    if hidden_states.device != self._split_device:
        hidden_states = hidden_states.to(self._split_device)
    if encoder_hidden_states.device != self._split_device:
        encoder_hidden_states = encoder_hidden_states.to(self._split_device)
    if temb.device != self._split_device:
        temb = temb.to(self._split_device)
    if image_rotary_emb is not None and image_rotary_emb[0].device != self._split_device:
        # is a tuple of tensors
        image_rotary_emb = tuple([t.to(self._split_device) for t in image_rotary_emb])
    return self._pre_gpu_split_forward(hidden_states, encoder_hidden_states, temb, image_rotary_emb, joint_attention_kwargs)


def split_gpu_single_block_forward(
    self,
    hidden_states: torch.FloatTensor,
    temb: torch.FloatTensor,
    image_rotary_emb=None,
    joint_attention_kwargs=None,
    **kwargs
):
    if hidden_states.device != self._split_device:
        hidden_states = hidden_states.to(device=self._split_device)
    if temb.device != self._split_device:
        temb = temb.to(device=self._split_device)
    if image_rotary_emb is not None and image_rotary_emb[0].device != self._split_device:
        # is a tuple of tensors
        image_rotary_emb = tuple([t.to(self._split_device) for t in image_rotary_emb])
    
    hidden_state_out = self._pre_gpu_split_forward(hidden_states, temb, image_rotary_emb, joint_attention_kwargs, **kwargs)
    if hasattr(self, "_split_output_device"):
        return hidden_state_out.to(self._split_output_device)
    return hidden_state_out


def add_model_gpu_splitter_to_flux(
    transformer: FluxTransformer2DModel,
    # ~ 5 billion for all other params
    other_module_params: Optional[int] = 5e9,
    # since they are not trainable, multiply by smaller number
    other_module_param_count_scale: Optional[float] = 0.3
):
    gpu_id_list = [i for i in range(torch.cuda.device_count())]
    
    # if len(gpu_id_list) > 2:
    #     raise ValueError("Cannot split to more than 2 GPUs currently.")
    other_module_params *= other_module_param_count_scale
    
    # since we are not tuning the 
    total_params = sum(p.numel() for p in transformer.parameters()) + other_module_params
    
    params_per_gpu = total_params / len(gpu_id_list)
    
    current_gpu_idx = 0
    # text encoders, vae, and some non block layers will all be on gpu 0
    current_gpu_params = other_module_params

    for double_block in transformer.transformer_blocks:
        device = torch.device(f"cuda:{current_gpu_idx}")
        double_block._pre_gpu_split_forward = double_block.forward
        double_block.forward = partial(
            split_gpu_double_block_forward, double_block)
        double_block._split_device = device
        # add the params to the current gpu
        current_gpu_params += sum(p.numel() for p in double_block.parameters())
        # if the current gpu params are greater than the params per gpu, move to next gpu
        if current_gpu_params > params_per_gpu:
            current_gpu_idx += 1
            current_gpu_params = 0
            if current_gpu_idx >= len(gpu_id_list):
                current_gpu_idx = gpu_id_list[-1]
        
    for single_block in transformer.single_transformer_blocks:
        device = torch.device(f"cuda:{current_gpu_idx}")
        single_block._pre_gpu_split_forward = single_block.forward
        single_block.forward = partial(
            split_gpu_single_block_forward, single_block)
        single_block._split_device = device
        # add the params to the current gpu
        current_gpu_params += sum(p.numel() for p in single_block.parameters())
        # if the current gpu params are greater than the params per gpu, move to next gpu
        if current_gpu_params > params_per_gpu:
            current_gpu_idx += 1
            current_gpu_params = 0
            if current_gpu_idx >= len(gpu_id_list):
                current_gpu_idx = gpu_id_list[-1]
    
    # add output device to last layer
    transformer.single_transformer_blocks[-1]._split_output_device = torch.device("cuda:0")
    
    transformer._pre_gpu_split_to = transformer.to
    transformer.to = partial(new_device_to, transformer)