Spaces:
Paused
Paused
File size: 7,107 Bytes
1c72248 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
# 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)
|