r""" Returns: A state dict containing just the LoRA parameters.
(unet: UNet2DConditionModel)
| 175 | |
| 176 | |
| 177 | def unet_lora_state_dict(unet: UNet2DConditionModel) -> Dict[str, torch.Tensor]: |
| 178 | r""" |
| 179 | Returns: |
| 180 | A state dict containing just the LoRA parameters. |
| 181 | """ |
| 182 | lora_state_dict = {} |
| 183 | |
| 184 | for name, module in unet.named_modules(): |
| 185 | if hasattr(module, "set_lora_layer"): |
| 186 | lora_layer = getattr(module, "lora_layer") |
| 187 | if lora_layer is not None: |
| 188 | current_lora_layer_sd = lora_layer.state_dict() |
| 189 | for lora_layer_matrix_name, lora_param in current_lora_layer_sd.items(): |
| 190 | # The matrix name can either be "down" or "up". |
| 191 | lora_state_dict[f"{name}.lora.{lora_layer_matrix_name}"] = lora_param |
| 192 | |
| 193 | return lora_state_dict |
| 194 | |
| 195 | |
| 196 | def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32): |
nothing calls this directly
no test coverage detected