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