ai-toolkit/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
from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, CombinedTimestepGuidanceTextProjEmbeddings


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)