How to load loRA weights?
Hi, I am trying to integrate a lora model into the stable diffusion model. I downloaded the checkpoints where loRAPath is the 'moxin_v1.0_lora_f16.ckpt' and unetPath is 'sd-v1.5.ckpt' from https://static.libnnc.org/sd-v1.5.ckpt
I noticed that the names of the weights for the two checkpoints are different, and so when I try to load the model I am doing something along the lines of (code below), where I manually change the name to match the key in the checkpoint.
I try to iterate through all the keys in the model and load the loRA weights where I can, and set them to zero if they are not in the lora checkpoints. My unet model is a LoRAUNet object.
However, when I run this script, the inference images do not seem to have any loRA effect. The results are the exact same as if I had not loaded in any loRA, so I am wondering if there's something I might be missing when trying to load loRA weights into the model?
Any advice is greatly appreciated! Thanks!
```
graph.openStore(loRAPath) { loraStore in
let loraKeys = Set(loraStore.keys)
graph.openStore(unetPath) { unetStore in
unetStore.read("unet", model: unet!, codec: codec) { name, _, _, _ in
if name.contains("lora_up") || name.contains("lora_down"){
// change the name to match the keys in the checkpoint
var loraName = name
if name.contains("lora_up"){
loraName = loraName.replacingOccurrences(of: "-lora_up", with: "")
loraName = loraName + "__up__"
} else if name.contains("lora_down"){
loraName = loraName.replacingOccurrences(of: "-lora_down", with: "")
loraName = loraName + "__down__"
}
if loraKeys.contains(loraName){
let original = graph.variable(Tensor<UseFloatingPoint>(from: loraStore.read(loraName)!))
return .final(original.rawValue)
}
else {
// replace the low rank matrix with all zeros
let lowRank = 16
if loraName.contains("__up__"){
loraName = loraName.replacingOccurrences(of: "__up__", with: "")
let value = graph.variable(Tensor<UseFloatingPoint>(from: unetStore.read(loraName, codec: codec)!)).toGPU(0)
var shape = value.shape
if shape.count == 4{
let upMatrix = graph.variable(Tensor<UseFloatingPoint>(.CPU, .NCHW(shape[0], lowRank, shape[2], shape[3]))).toGPU(0)
upMatrix.full(0)
return .final(upMatrix.rawValue.toCPU())
} else if shape.count == 3{
let upMatrix = graph.variable(Tensor<UseFloatingPoint>(.CPU, .CHW(shape[0], lowRank, shape[2]))).toGPU(0)
upMatrix.full(0)
return .final(upMatrix.rawValue.toCPU())
} else if shape.count == 2{
let upMatrix = graph.variable(Tensor<UseFloatingPoint>(.CPU, .NC(shape[0], lowRank))).toGPU(0)
upMatrix.full(0)
return .final(upMatrix.rawValue.toCPU())
} else {
value.full(0)
print("\n shape :: \(shape.count)")
return .final(value.rawValue.toCPU())
}
} else{
loraName = loraName.replacingOccurrences(of: "__down__", with: "")
let value = graph.variable(Tensor<UseFloatingPoint>(from: unetStore.read(loraName, codec: codec)!)).toGPU(0)
var shape = value.shape
if shape.count == 4{
let downMatrix = graph.variable(Tensor<UseFloatingPoint>(.CPU, .NCHW(lowRank, shape[1], shape[2], shape[3]))).toGPU(0)
downMatrix.full(0)
return .final(downMatrix.rawValue.toCPU())
} else if shape.count == 3{
let downMatrix = graph.variable(Tensor<UseFloatingPoint>(.CPU, .CHW(lowRank, shape[1], shape[2]))).toGPU(0)
downMatrix.full(0)
return .final(downMatrix.rawValue.toCPU())
} else if shape.count == 2{
let downMatrix = graph.variable(Tensor<UseFloatingPoint>(.CPU, .NC(lowRank, shape[1]))).toGPU(0)
downMatrix.full(0)
return .final(downMatrix.rawValue.toCPU())
} else {
value.full(0)
return .final(value.rawValue.toCPU())
}
}
}
}
else {
return .continue(name)
}
}
```
关闭于 2024-04-08 16 条评论