Handle v2 models LoRA.

Author: liuliu

examples/lora/main.swift

@@ -110,7 +110,7 @@ public final class SafeTensors {
  }
 }
 
-let filename = "/home/liu/workspace/swift-diffusion/openjourneyLora_v1.safetensors"
+let filename = "/home/liu/workspace/swift-diffusion/to8sHighKeyLORASD15SD2_sd21768.safetensors"
 /*
 let archive = Archive(url: URL(fileURLWithPath: filename), accessMode: .read)!
 let entry = archive["archive/data.pkl"]!
@@ -172,7 +172,7 @@ for i in stride(from: 0, to: unetMap.count, by: 2) {
  let parts = unetMap[i].components(separatedBy: ".")
  unetMap[i] = parts[2..<(parts.count - 1)].joined(separator: "_") + "." + parts[parts.count - 1]
 }
-var textModelMap = try jsonDecoder.decode([String].self, from: Data(contentsOf: URL(fileURLWithPath: "/home/liu/workspace/swift-diffusion/text_model.json")))
+var textModelMap = try jsonDecoder.decode([String].self, from: Data(contentsOf: URL(fileURLWithPath: "/home/liu/workspace/swift-diffusion/open_clip_text_model_1.json")))
 for i in stride(from: 0, to: textModelMap.count, by: 2) {
  let parts = textModelMap[i].components(separatedBy: ".")
  textModelMap[i] = parts[2..<(parts.count - 1)].joined(separator: "_") + "." + parts[parts.count - 1]

examples/open_clip/main.swift

@@ -159,12 +159,16 @@ let vocab = state_dict["token_embedding.weight"]
 let pos = state_dict["positional_embedding"]
 tokenEmbed.parameters.copy(from: try! Tensor<Float>(numpy: vocab.numpy()))
 positionEmbed.parameters.copy(from: try! Tensor<Float>(numpy: pos.numpy()))
+print("\"token_embedding.weight\", \"\(tokenEmbed.parameters.name)\"")
+print("\"positional_embedding\", \"\(positionEmbed.parameters.name)\"")
 
 for i in 0..<23 {
  let layer_norm_1_weight = state_dict["transformer.resblocks.\(i).ln_1.weight"].numpy()
  let layer_norm_1_bias = state_dict["transformer.resblocks.\(i).ln_1.bias"].numpy()
  layerNorm1s[i].parameters(for: .weight).copy(from: try! Tensor<Float>(numpy: layer_norm_1_weight))
  layerNorm1s[i].parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: layer_norm_1_bias))
+ print("\"transformer.resblocks.\(i).ln_1.weight\", \"\(layerNorm1s[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).ln_1.bias\", \"\(layerNorm1s[i].parameters(for: .bias).name)\"")
 
  let in_proj_weight = state_dict["transformer.resblocks.\(i).attn.in_proj_weight"].type(
  torch.float
@@ -174,42 +178,58 @@ for i in 0..<23 {
  toqueries[i].parameters(for: .weight).copy(
  from: try! Tensor<Float>(numpy: in_proj_weight[..<(1024), ...]))
  toqueries[i].parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: in_proj_bias[..<(1024)]))
+ print("\"transformer.resblocks.\(i).attn.in_proj_weight\", \"\(toqueries[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).attn.in_proj_bias\", \"\(toqueries[i].parameters(for: .bias).name)\"")
  tokeys[i].parameters(for: .weight).copy(
  from: try! Tensor<Float>(numpy: in_proj_weight[(1024)..<(2 * 1024), ...]))
  tokeys[i].parameters(for: .bias).copy(
  from: try! Tensor<Float>(numpy: in_proj_bias[(1024)..<(2 * 1024)]))
+ print("\"transformer.resblocks.\(i).attn.in_proj_weight\", \"\(tokeys[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).attn.in_proj_bias\", \"\(tokeys[i].parameters(for: .bias).name)\"")
  tovalues[i].parameters(for: .weight).copy(
  from: try! Tensor<Float>(numpy: in_proj_weight[(2 * 1024)..., ...]))
  tovalues[i].parameters(for: .bias).copy(
  from: try! Tensor<Float>(numpy: in_proj_bias[(2 * 1024)...]))
+ print("\"transformer.resblocks.\(i).attn.in_proj_weight\", \"\(tovalues[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).attn.in_proj_bias\", \"\(tovalues[i].parameters(for: .bias).name)\"")
 
  let out_proj_weight = state_dict["transformer.resblocks.\(i).attn.out_proj.weight"]
  .numpy()
  let out_proj_bias = state_dict["transformer.resblocks.\(i).attn.out_proj.bias"].numpy()
  unifyheads[i].parameters(for: .weight).copy(from: try! Tensor<Float>(numpy: out_proj_weight))
  unifyheads[i].parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: out_proj_bias))
+ print("\"transformer.resblocks.\(i).attn.out_proj.weight\", \"\(unifyheads[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).attn.out_proj.bias\", \"\(unifyheads[i].parameters(for: .bias).name)\"")
 
  let layer_norm_2_weight = state_dict["transformer.resblocks.\(i).ln_2.weight"].numpy()
  let layer_norm_2_bias = state_dict["transformer.resblocks.\(i).ln_2.bias"].numpy()
  layerNorm2s[i].parameters(for: .weight).copy(from: try! Tensor<Float>(numpy: layer_norm_2_weight))
  layerNorm2s[i].parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: layer_norm_2_bias))
+ print("\"transformer.resblocks.\(i).ln_2.weight\", \"\(layerNorm2s[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).ln_2.bias\", \"\(layerNorm2s[i].parameters(for: .bias).name)\"")
 
  let fc1_weight = state_dict["transformer.resblocks.\(i).mlp.c_fc.weight"].numpy()
  let fc1_bias = state_dict["transformer.resblocks.\(i).mlp.c_fc.bias"].numpy()
  fc1s[i].parameters(for: .weight).copy(from: try! Tensor<Float>(numpy: fc1_weight))
  fc1s[i].parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: fc1_bias))
+ print("\"transformer.resblocks.\(i).mlp.c_fc.weight\", \"\(fc1s[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).mlp.c_fc.bias\", \"\(fc1s[i].parameters(for: .bias).name)\"")
 
  let fc2_weight = state_dict["transformer.resblocks.\(i).mlp.c_proj.weight"].numpy()
  let fc2_bias = state_dict["transformer.resblocks.\(i).mlp.c_proj.bias"].numpy()
  fc2s[i].parameters(for: .weight).copy(from: try! Tensor<Float>(numpy: fc2_weight))
  fc2s[i].parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: fc2_bias))
+ print("\"transformer.resblocks.\(i).mlp.c_proj.weight\", \"\(fc2s[i].parameters(for: .weight).name)\"")
+ print("\"transformer.resblocks.\(i).mlp.c_proj.bias\", \"\(fc2s[i].parameters(for: .bias).name)\"")
 }
 
 let final_layer_norm_weight = state_dict["ln_final.weight"].numpy()
 let final_layer_norm_bias = state_dict["ln_final.bias"].numpy()
 finalLayerNorm.parameters(for: .weight).copy(
  from: try! Tensor<Float>(numpy: final_layer_norm_weight))
 finalLayerNorm.parameters(for: .bias).copy(from: try! Tensor<Float>(numpy: final_layer_norm_bias))
+print("\"ln_final.weight\", \"\(finalLayerNorm.parameters(for: .weight).name)\"")
+print("\"ln_final.bias\", \"\(finalLayerNorm.parameters(for: .bias).name)\"")
 
 let c = textModel(inputs: tokensTensor, positionTensor, casualAttentionMask)[0].as(of: Float.self)
 for i in 0..<6 {