content: clip_param_keys = {
      'positional_embedding',
      'text_projection',
      'visual.class_embedding',
      'visual.positional_embedding',
      'visual.conv1.weight',
      'visual.ln_pre.weight',
      'visual.ln_pre.bias',
      'visual.transformer.resblocks.0.attn.in_proj_weight',
      'visual.transformer.resblocks.0.attn.in_proj_bias',
      'visual.transformer.resblocks.0.attn.out_proj.weight',
      'visual.transformer.resblocks.0.attn.out_proj.bias',
      'visual.transformer.resblocks.0.ln_1.weight',
      'visual.transformer.resblocks.0.ln_1.bias',
      'visual.transformer.resblocks.0.mlp.c_fc.weight',
      'visual.transformer.resblocks.0.mlp.c_fc.bias',
      'visual.transformer.resblocks.0.mlp.c_proj.weight',
      'visual.transformer.resblocks.0.mlp.c_proj.bias',
      'visual.transformer.resblocks.0.ln_2.weight',
      'visual.transformer.resblocks.0.ln_2.bias',
      'visual.transformer.resblocks.1.attn.in_proj_weight',
      'visual.transformer.resblocks.1.attn.in_proj_bias',
      'visual.transformer.resblocks.1.attn.out_proj.weight',
      'visual.transformer.resblocks.1.attn.out_proj.bias',
      'visual.transformer.resblocks.1.ln_1.weight',
      'visual.transformer.resblocks.1.ln_1.bias',
      'visual.transformer.resblocks.1.mlp.c_fc.weight',
      'visual.transformer.resblocks.1.mlp.c_fc.bias',
      'visual.transformer.resblocks.1.mlp.c_proj.weight',
      'visual.transformer.resblocks.1.mlp.c_proj.bias',
      'visual.transformer.resblocks.1.ln_2.weight',
      'visual.transformer.resblocks.1.ln_2.bias',
      'visual.transformer.resblocks.2.attn.in_proj_weight',
      'visual.transformer.resblocks.2.attn.in_proj_bias',
      'visual.transformer.resblocks.2.attn.out_proj.weight',
      'visual.transformer.resblocks.2.attn.out_proj.bias',
      'visual.transformer.resblocks.2.ln_1.weight',
      'visual.transformer.resblocks.2.ln_1.bias',
      'visual.transformer.resblocks.2.mlp.c_fc.weight',
      'visual.transformer.resblocks.2.mlp.c_fc.bias',
      'visual.transformer.resblocks.2.mlp.c_proj.weight',
      'visual.transformer.resblocks.2.mlp.c_proj.bias',
      'visual.transformer.resblocks.2.ln_2.weight',
      'visual.transformer.resblocks.2.ln_2.bias',
      'visual.transformer.resblocks.3.attn.in_proj_weight',
      'visual.transformer.resblocks.3.attn.in_proj_bias',
      'visual.transformer.resblocks.3.attn.out_proj.weight',
      'visual.transformer.resblocks.3.attn.out_proj.bias',
      'visual.transformer.resblocks.3.ln_1.weight',
      'visual.transformer.resblocks.3.ln_1.bias',
      'visual.transformer.resblocks.3.mlp.c_fc.weight',
      'visual.transformer.resblocks.3.mlp.c_fc.bias',
      'visual.transformer.resblocks.3.mlp.c_proj.weight',
      'visual.transformer.resblocks.3.mlp.c_proj.bias',
      'visual.transformer.resblocks.3.ln_2.weight',
      'visual.transformer.resblocks.3.ln_2.bias',
      'visual.transformer.resblocks.4.attn.in_proj_weight',
      'visual.transformer.resblocks.4.attn.in_proj_bias',
      'visual.transformer.resblocks.4.attn.out_proj.weight',
      'visual.transformer.resblocks.4.attn.out_proj.bias',
      'visual.transformer.resblocks.4.ln_1.weight',
      'visual.transformer.resblocks.4.ln_1.bias',
      'visual.transformer.resblocks.4.mlp.c_fc.weight',
      'visual.transformer.resblocks.4.mlp.c_fc.bias',
      'visual.transformer.resblocks.4.mlp.c_proj.weight',
      'visual.transformer.resblocks.4.mlp.c_proj.bias',
      'visual.transformer.resblocks.4.ln_2.weight',
      'visual.transformer.resblocks.4.ln_2.bias',
      'visual.transformer.resblocks.5.attn.in_proj_weight',
      'visual.transformer.resblocks.5.attn.in_proj_bias',
      'visual.transformer.resblocks.5.attn.out_proj.weight',
      'visual.transformer.resblocks.5.attn.out_proj.bias',
      'visual.transformer.resblocks.5.ln_1.weight',
      'visual.transformer.resblocks.5.ln_1.bias',
      'visual.transformer.resblocks.5.mlp.c_fc.weight',
      'visual.transformer.resblocks.5.mlp.c_fc.bias',
      'visual.transformer.resblocks.5.mlp.c_proj.weight',
      'visual.transformer.resblocks.5.mlp.c_proj.bias',
      'visual.transformer.resblocks.5.ln_2.weight',
      'visual.transformer.resblocks.5.ln_2.bias',
      'visual.transformer.resblocks.6.attn.in_proj_weight',
      'visual.transformer.resblocks.6.attn.in_proj_bias',
      'visual.transformer.resblocks.6.attn.out_proj.weight',
      'visual.transformer.resblocks.6.attn.out_proj.bias',
      'visual.transformer.resblocks.6.ln_1.weight',
      'visual.transformer.resblocks.6.ln_1.bias',
      'visual.transformer.resblocks.6.mlp.c_fc.weight',
      'visual.transformer.resblocks.6.mlp.c_fc.bias',
      'visual.transformer.resblocks.6.mlp.c_proj.weight',
      'visual.transformer.resblocks.6.mlp.c_proj.bias',
      'visual.transformer.resblocks.6.ln_2.weight',
      'visual.transformer.resblocks.6.ln_2.bias',
      'visual.transformer.resblocks.7.attn.in_proj_weight',
      'visual.transformer.resblocks.7.attn.in_proj_bias',
      'visual.transformer.resblocks.7.attn.out_proj.weight',
      'visual.transformer.resblocks.7.attn.out_proj.bias',
      'visual.transformer.resblocks.7.ln_1.weight',
      'visual.transformer.resblocks.7.ln_1.bias',
      'visual.transformer.resblocks.7.mlp.c_fc.weight',
      'visual.transformer.resblocks.7.mlp.c_fc.bias',
      'visual.transformer.resblocks.7.mlp.c_proj.weight',
      'visual.transformer.resblocks.7.mlp.c_proj.bias',
      'visual.transformer.resblocks.7.ln_2.weight',
      'visual.transformer.resblocks.7.ln_2.bias',
      'visual.transformer.resblocks.8.attn.in_proj_weight',
      'visual.transformer.resblocks.8.attn.in_proj_bias',
      'visual.transformer.resblocks.8.attn.out_proj.weight',
      'visual.transformer.resblocks.8.attn.out_proj.bias',
      'visual.transformer.resblocks.8.ln_1.weight',
      'visual.transformer.resblocks.8.ln_1.bias',
      'visual.transformer.resblocks.8.mlp.c_fc.weight',
      'visual.transformer.resblocks.8.mlp.c_fc.bias',
      'visual.transformer.resblocks.8.mlp.c_proj.weight',
      'visual.transformer.resblocks.8.mlp.c_proj.bias',
      'visual.transformer.resblocks.8.ln_2.weight',
      'visual.transformer.resblocks.8.ln_2.bias',
      'visual.transformer.resblocks.9.attn.in_proj_weight',
      'visual.transformer.resblocks.9.attn.in_proj_bias',
      'visual.transformer.resblocks.9.attn.out_proj.weight',
      'visual.transformer.resblocks.9.attn.out_proj.bias',
      'visual.transformer.resblocks.9.ln_1.weight',
      'visual.transformer.resblocks.9.ln_1.bias',
      'visual.transformer.resblocks.9.mlp.c_fc.weight',
      'visual.transformer.resblocks.9.mlp.c_fc.bias',
      'visual.transformer.resblocks.9.mlp.c_proj.weight',
      'visual.transformer.resblocks.9.mlp.c_proj.bias',
      'visual.transformer.resblocks.9.ln_2.weight',
      'visual.transformer.resblocks.9.ln_2.bias',
      'visual.transformer.resblocks.10.attn.in_proj_weight',
      'visual.transformer.resblocks.10.attn.in_proj_bias',
      'visual.transformer.resblocks.10.attn.out_proj.weight',
      'visual.transformer.resblocks.10.attn.out_proj.bias',
      'visual.transformer.resblocks.10.ln_1.weight',
      'visual.transformer.resblocks.10.ln_1.bias',
      'visual.transformer.resblocks.10.mlp.c_fc.weight',
      'visual.transformer.resblocks.10.mlp.c_fc.bias',
      'visual.transformer.resblocks.10.mlp.c_proj.weight',
      'visual.transformer.resblocks.10.mlp.c_proj.bias',
      'visual.transformer.resblocks.10.ln_2.weight',
      'visual.transformer.resblocks.10.ln_2.bias',
      'visual.transformer.resblocks.11.attn.in_proj_weight',
      'visual.transformer.resblocks.11.attn.in_proj_bias',
      'visual.transformer.resblocks.11.attn.out_proj.weight',
      'visual.transformer.resblocks.11.attn.out_proj.bias',
      'visual.transformer.resblocks.11.ln_1.weight',
      'visual.transformer.resblocks.11.ln_1.bias',
      'visual.transformer.resblocks.11.mlp.c_fc.weight',
      'visual.transformer.resblocks.11.mlp.c_fc.bias',
      'visual.transformer.resblocks.11.mlp.c_proj.weight',
      'visual.transformer.resblocks.11.mlp.c_proj.bias',
      'visual.transformer.resblocks.11.ln_2.weight',
      'visual.transformer.resblocks.11.ln_2.bias',
      'visual.ln_post.weight',
      'visual.ln_post.bias',
      'visual_ln_post.weight',
      'visual_ln_post.bias',
      'transformer.resblocks.0.attn.in_proj_weight',
      'transformer.resblocks.0.attn.in_proj_bias',
      'transformer.resblocks.0.attn.out_proj.weight',
      'transformer.resblocks.0.attn.out_proj.bias',
      'transformer.resblocks.0.ln_1.weight',
      'transformer.resblocks.0.ln_1.bias',
      'transformer.resblocks.0.mlp.c_fc.weight',
      'transformer.resblocks.0.mlp.c_fc.bias',
      'transformer.resblocks.0.mlp.c_proj.weight',
      'transformer.resblocks.0.mlp.c_proj.bias',
      'transformer.resblocks.0.ln_2.weight',
      'transformer.resblocks.0.ln_2.bias',
      'transformer.resblocks.1.attn.in_proj_weight',
      'transformer.resblocks.1.attn.in_proj_bias',
      'transformer.resblocks.1.attn.out_proj.weight',
      'transformer.resblocks.1.attn.out_proj.bias',
      'transformer.resblocks.1.ln_1.weight',
      'transformer.resblocks.1.ln_1.bias',
      'transformer.resblocks.1.mlp.c_fc.weight',
      'transformer.resblocks.1.mlp.c_fc.bias',
      'transformer.resblocks.1.mlp.c_proj.weight',
      'transformer.resblocks.1.mlp.c_proj.bias',
      'transformer.resblocks.1.ln_2.weight',
      'transformer.resblocks.1.ln_2.bias',
      'transformer.resblocks.2.attn.in_proj_weight',
      'transformer.resblocks.2.attn.in_proj_bias',
      'transformer.resblocks.2.attn.out_proj.weight',
      'transformer.resblocks.2.attn.out_proj.bias',
      'transformer.resblocks.2.ln_1.weight',
      'transformer.resblocks.2.ln_1.bias',
      'transformer.resblocks.2.mlp.c_fc.weight',
      'transformer.resblocks.2.mlp.c_fc.bias',
      'transformer.resblocks.2.mlp.c_proj.weight',
      'transformer.resblocks.2.mlp.c_proj.bias',
      'transformer.resblocks.2.ln_2.weight',
      'transformer.resblocks.2.ln_2.bias',
      'transformer.resblocks.3.attn.in_proj_weight',
      'transformer.resblocks.3.attn.in_proj_bias',
      'transformer.resblocks.3.attn.out_proj.weight',
      'transformer.resblocks.3.attn.out_proj.bias',
      'transformer.resblocks.3.ln_1.weight',
      'transformer.resblocks.3.ln_1.bias',
      'transformer.resblocks.3.mlp.c_fc.weight',
      'transformer.resblocks.3.mlp.c_fc.bias',
      'transformer.resblocks.3.mlp.c_proj.weight',
      'transformer.resblocks.3.mlp.c_proj.bias',
      'transformer.resblocks.3.ln_2.weight',
      'transformer.resblocks.3.ln_2.bias',
      'transformer.resblocks.4.attn.in_proj_weight',
      'transformer.resblocks.4.attn.in_proj_bias',
      'transformer.resblocks.4.attn.out_proj.weight',
      'transformer.resblocks.4.attn.out_proj.bias',
      'transformer.resblocks.4.ln_1.weight',
      'transformer.resblocks.4.ln_1.bias',
      'transformer.resblocks.4.mlp.c_fc.weight',
      'transformer.resblocks.4.mlp.c_fc.bias',
      'transformer.resblocks.4.mlp.c_proj.weight',
      'transformer.resblocks.4.mlp.c_proj.bias',
      'transformer.resblocks.4.ln_2.weight',
      'transformer.resblocks.4.ln_2.bias',
      'transformer.resblocks.5.attn.in_proj_weight',
      'transformer.resblocks.5.attn.in_proj_bias',
      'transformer.resblocks.5.attn.out_proj.weight',
      'transformer.resblocks.5.attn.out_proj.bias',
      'transformer.resblocks.5.ln_1.weight',
      'transformer.resblocks.5.ln_1.bias',
      'transformer.resblocks.5.mlp.c_fc.weight',
      'transformer.resblocks.5.mlp.c_fc.bias',
      'transformer.resblocks.5.mlp.c_proj.weight',
      'transformer.resblocks.5.mlp.c_proj.bias',
      'transformer.resblocks.5.ln_2.weight',
      'transformer.resblocks.5.ln_2.bias',
      'transformer.resblocks.6.attn.in_proj_weight',
      'transformer.resblocks.6.attn.in_proj_bias',
      'transformer.resblocks.6.attn.out_proj.weight',
      'transformer.resblocks.6.attn.out_proj.bias',
      'transformer.resblocks.6.ln_1.weight',
      'transformer.resblocks.6.ln_1.bias',
      'transformer.resblocks.6.mlp.c_fc.weight',
      'transformer.resblocks.6.mlp.c_fc.bias',
      'transformer.resblocks.6.mlp.c_proj.weight',
      'transformer.resblocks.6.mlp.c_proj.bias',
      'transformer.resblocks.6.ln_2.weight',
      'transformer.resblocks.6.ln_2.bias',
      'transformer.resblocks.7.attn.in_proj_weight',
      'transformer.resblocks.7.attn.in_proj_bias',
      'transformer.resblocks.7.attn.out_proj.weight',
      'transformer.resblocks.7.attn.out_proj.bias',
      'transformer.resblocks.7.ln_1.weight',
      'transformer.resblocks.7.ln_1.bias',
      'transformer.resblocks.7.mlp.c_fc.weight',
      'transformer.resblocks.7.mlp.c_fc.bias',
      'transformer.resblocks.7.mlp.c_proj.weight',
      'transformer.resblocks.7.mlp.c_proj.bias',
      'transformer.resblocks.7.ln_2.weight',
      'transformer.resblocks.7.ln_2.bias',
      'transformer.resblocks.8.attn.in_proj_weight',
      'transformer.resblocks.8.attn.in_proj_bias',
      'transformer.resblocks.8.attn.out_proj.weight',
      'transformer.resblocks.8.attn.out_proj.bias',
      'transformer.resblocks.8.ln_1.weight',
      'transformer.resblocks.8.ln_1.bias',
      'transformer.resblocks.8.mlp.c_fc.weight',
      'transformer.resblocks.8.mlp.c_fc.bias',
      'transformer.resblocks.8.mlp.c_proj.weight',
      'transformer.resblocks.8.mlp.c_proj.bias',
      'transformer.resblocks.8.ln_2.weight',
      'transformer.resblocks.8.ln_2.bias',
      'transformer.resblocks.9.attn.in_proj_weight',
      'transformer.resblocks.9.attn.in_proj_bias',
      'transformer.resblocks.9.attn.out_proj.weight',
      'transformer.resblocks.9.attn.out_proj.bias',
      'transformer.resblocks.9.ln_1.weight',
      'transformer.resblocks.9.ln_1.bias',
      'transformer.resblocks.9.mlp.c_fc.weight',
      'transformer.resblocks.9.mlp.c_fc.bias',
      'transformer.resblocks.9.mlp.c_proj.weight',
      'transformer.resblocks.9.mlp.c_proj.bias',
      'transformer.resblocks.9.ln_2.weight',
      'transformer.resblocks.9.ln_2.bias',
      'transformer.resblocks.10.attn.in_proj_weight',
      'transformer.resblocks.10.attn.in_proj_bias',
      'transformer.resblocks.10.attn.out_proj.weight',
      'transformer.resblocks.10.attn.out_proj.bias',
      'transformer.resblocks.10.ln_1.weight',
      'transformer.resblocks.10.ln_1.bias',
      'transformer.resblocks.10.mlp.c_fc.weight',
      'transformer.resblocks.10.mlp.c_fc.bias',
      'transformer.resblocks.10.mlp.c_proj.weight',
      'transformer.resblocks.10.mlp.c_proj.bias',
      'transformer.resblocks.10.ln_2.weight',
      'transformer.resblocks.10.ln_2.bias',
      'transformer.resblocks.11.attn.in_proj_weight',
      'transformer.resblocks.11.attn.in_proj_bias',
      'transformer.resblocks.11.attn.out_proj.weight',
      'transformer.resblocks.11.attn.out_proj.bias',
      'transformer.resblocks.11.ln_1.weight',
      'transformer.resblocks.11.ln_1.bias',
      'transformer.resblocks.11.mlp.c_fc.weight',
      'transformer.resblocks.11.mlp.c_fc.bias',
      'transformer.resblocks.11.mlp.c_proj.weight',
      'transformer.resblocks.11.mlp.c_proj.bias',
      'transformer.resblocks.11.ln_2.weight',
      'transformer.resblocks.11.ln_2.bias',
      'token_embedding.weight',
      'visual.proj',
      'visual_proj',
      'ln_final.weight',
      'ln_final.bias',
  }
  
  
  def gradually_freeze_by_layer(model, global_step, interval, unfreeze_type="both"):
      try:
          visual_n_layers = model.clip.visual.transformer.layers
      except:
          visual_n_layers = model.visual.transformer.layers
  
      try:
          text_n_layers = model.clip.transformer.layers
      except:
          text_n_layers = model.transformer.layers
  
      if interval <= 0:
          return
  
      layer_to_unfreeze = global_step // interval
      if (
          visual_n_layers + 1 < layer_to_unfreeze
          and text_n_layers + 1 < layer_to_unfreeze
      ):  # Do nothing
          return
      embeddings = [
          "token_embedding",
          "visual.positional_embedding" "visual.class_embedding",
          "positional_embedding",
      ]
      visual_first_keys = [
          "visual.conv1.weight",
          "visual.ln_pre",
      ]
      text_first_keys = []
      visual_last_keys = ["visual.ln_post", "visual_ln_post" "visual_proj", "visual.proj"]
      text_last_keys = ["ln_final"]
  
      def set_unfreeze(n, p, l):
          if any(x in n for x in l):
              p.requires_grad = True
  
      for n, p in model.named_parameters():
          if layer_to_unfreeze == 0:  # At the beginning of training
              set_unfreeze(n, p, embeddings)  # Unfreeze embeddings
              if unfreeze_type in ["both", "visual"]:
                  set_unfreeze(n, p, visual_last_keys)
              if unfreeze_type in ["both", "text"]:  # Unfreeze text last layers
                  set_unfreeze(n, p, text_last_keys)
              continue
  
          if unfreeze_type in ["both", "visual"]:
              if layer_to_unfreeze <= visual_n_layers:  # Unfreeze n-th visual layer
                  param_prefix = f"visual.transformer.resblocks.{visual_n_layers - layer_to_unfreeze}"
                  set_unfreeze(n, p, [param_prefix])
              else:  # Unfreeze first visual layers
                  set_unfreeze(n, p, visual_first_keys)
          if unfreeze_type in ["both", "text"]:
              if layer_to_unfreeze <= text_n_layers:  # Unfreeze n-th text layer
                  param_prefix = (
                      f"transformer.resblocks.{visual_n_layers - layer_to_unfreeze}"
                  )
                  set_unfreeze(n, p, [param_prefix])
              else:  # Unfreeze first visual layers
                  set_unfreeze(n, p, text_first_keys)
  