[diffusion] lora: fix LoRA dtype handling and weight attribute access for z-image model (#14543)
Co-authored-by: niehen6174 <nihen6174@gmail.com>
This commit is contained in:
co-authored by
niehen6174
parent
e12c6b320f
commit
80122e4f4c
@@ -59,6 +59,14 @@ class BaseLayerWithLoRA(nn.Module):
|
|||||||
self.lora_A = None
|
self.lora_A = None
|
||||||
self.lora_B = None
|
self.lora_B = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def weight(self):
|
||||||
|
return self.base_layer.weight
|
||||||
|
|
||||||
|
@property
|
||||||
|
def bias(self):
|
||||||
|
return getattr(self.base_layer, "bias", None)
|
||||||
|
|
||||||
@torch.compile()
|
@torch.compile()
|
||||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
lora_A = self.lora_A
|
lora_A = self.lora_A
|
||||||
@@ -79,7 +87,7 @@ class BaseLayerWithLoRA(nn.Module):
|
|||||||
return out + delta, output_bias
|
return out + delta, output_bias
|
||||||
else:
|
else:
|
||||||
out, output_bias = self.base_layer(x)
|
out, output_bias = self.base_layer(x)
|
||||||
return out.to(x), output_bias
|
return out, output_bias
|
||||||
|
|
||||||
def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
|
def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
|
||||||
return A
|
return A
|
||||||
|
|||||||
Reference in New Issue
Block a user