code: def generic_injection(module, dtype=None, enable_cuda_graph=True):
      def replace_attn(child, policy):
          policy_attn = policy.attention(child)
          if policy_attn is None:
              return child
          if len(policy_attn) == 5:
              qkvw, attn_ow, attn_ob, hidden_size, heads = policy_attn
          else:
              (
                  qw,
                  kw,
                  vw,
                  attn_ow,
                  attn_ob,
                  hidden_size,
                  heads,
              ) = policy_attn
  
          config = transformer_inference.DeepSpeedInferenceConfig(
              hidden_size=hidden_size,
              heads=heads,
              dtype=dtype,
              triangular_masking=False,
              max_out_tokens=4096,
          )
          attn_module = DeepSpeedDiffusersAttention(config)
  
          def transpose(data):
              data = data.contiguous()
              data.reshape(-1).copy_(
                  data.transpose(-1, -2).contiguous().reshape(-1)
              )
              data = data.reshape(data.shape[-1], data.shape[-2])
              data.to(get_accelerator().current_device_name())
              return data
  
          if len(policy_attn) == 5:
              attn_module.attn_qkvw.data = transpose(qkvw.data)
          else:
              attn_module.attn_qkvw = None
              attn_module.attn_qw.data = transpose(qw.data)
              attn_module.attn_kw.data = transpose(kw.data)
              attn_module.attn_vw.data = transpose(vw.data)
  
          attn_module.attn_qkvb = None
          attn_module.attn_ow.data = transpose(attn_ow.data)
          attn_module.attn_ob.data.copy_(
              attn_ob.data.to(get_accelerator().current_device_name())
          )
          return attn_module
  
      def replace_attn_block(child, policy):
          config = Diffusers2DTransformerConfig()
          return DeepSpeedDiffusersTransformerBlock(child, config)
  
      if isinstance(module, torch.nn.Module):
          pass
      else:
          if dtype not in [torch.float16, torch.half]:
              raise ValueError(
                  "Generic injection only supported with FP16"
              )
  
          try:
              import diffusers
  
              if hasattr(diffusers.models.attention, "CrossAttention"):
                  cross_attention = (
                      diffusers.models.attention.CrossAttention
                  )
              else:
                  cross_attention = (
                      diffusers.models.attention_processor.Attention
                  )
              attention_block = (
                  diffusers.models.attention.BasicTransformerBlock
              )
              new_policies = {
                  cross_attention: replace_attn,
                  attention_block: replace_attn_block,
              }
          except ImportError:
              new_policies = {}
  
          # replace_transformer_layer(None,
          #                          module.text_encoder,
          #                          training=False,
          #                          replace_with_kernel_inject=True,
          #                          triangular_masking=True,
          #                          max_out_tokens=8192)
          from ..model_implementations.transformers.clip_encoder import (
              DSClipEncoder,
          )
  
          cg_encoder = DSClipEncoder(
              module.text_encoder, enable_cuda_graph=enable_cuda_graph
          )
          setattr(module, "text_encoder", cg_encoder)
          for name in module.__dict__.keys():
              sub_module = getattr(module, name)
              policy = _module_match(sub_module)
  
              if policy is not None:
  
                  def _replace_module(module, policy):
                      for name, child in module.named_children():
                          _replace_module(child, policy)
                          if child.__class__ in new_policies:
                              replaced_module = new_policies[
                                  child.__class__
                              ](child, policy)
                              setattr(module, name, replaced_module)
  
                  _replace_module(sub_module, policy)
                  new_module = policy.apply(
                      sub_module, enable_cuda_graph=enable_cuda_graph
                  )
                  print(
                      f"**** found and replaced {name} w. {type(new_module)}"
                  )
                  setattr(module, name, new_module)
  
imports: import os
  import torch
  import tqdm
  import deepspeed
  import deepspeed.ops.transformer
  from deepspeed.ops.transformer.inference.diffusers_attention import DeepSpeedDiffusersAttention
  from deepspeed.ops.transformer.inference.diffusers_transformer_block import DeepSpeedDiffusersTransformerBlock
  from deepspeed.ops.transformer.inference.diffusers_2d_transformer import Diffusers2DTransformerConfig
  from deepspeed.accelerator import get_accelerator
  from .replace_policy import replace_policies, generic_policies
  from .auto_tp import AutoTP, ReplaceWithTensorSlicing, Loading
  from deepspeed import comm
  from deepspeed.module_inject.tp_shard import set_num_kv_heads
  from .load_checkpoint import load_model_with_checkpoint
  import time
  from .utils import policy_to_ds_container
  import gc
  from ..pipe import PipelineModule
  import re
label_desc: Use image node_id 3 for calling a global function with example usage: generic_injection(module, dtype, enable_cuda_graph) and returns: attn_module, DeepSpeedDiffusersTransformerBlock, child, data