ITADN

How to load loRA weights?

#56Closeddavidw0311 创建于 2024-03-19
D
davidw0311commented
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 条评论