# 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)