Merge branch 'sd3' into feat-support-lokr-loha

docs/anima_train_network.md

@@ -652,4 +652,4 @@ The following metadata is saved in the LoRA model file:
 * `ss_sigmoid_scale`
 * `ss_discrete_flow_shift`
 
-</details>
+</details>

networks/lora_anima.py

@@ -840,4 +840,4 @@ class LoRANetwork(torch.nn.Module):
             scalednorm = updown.norm() * ratio
             norms.append(scalednorm.item())
 
-        return keys_scaled, sum(norms) / len(norms), max(norms)
+        return keys_scaled, sum(norms) / len(norms), max(norms)

networks/lora_flux.py

@@ -141,10 +141,13 @@ class LoRAModule(torch.nn.Module):
             # rank dropout
             if self.rank_dropout is not None and self.training:
                 mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout
-                if len(lx.size()) == 3:
-                    mask = mask.unsqueeze(1)  # for Text Encoder
-                elif len(lx.size()) == 4:
-                    mask = mask.unsqueeze(-1).unsqueeze(-1)  # for Conv2d
+                if isinstance(self.lora_down, torch.nn.Conv2d):
+                    # Conv2d: lora_dim is at dim 1 → [B, dim, 1, 1]
+                    mask = mask.unsqueeze(-1).unsqueeze(-1)
+                else:
+                    # Linear: lora_dim is at last dim → [B, 1, ..., 1, dim]
+                    for _ in range(len(lx.size()) - 2):
+                        mask = mask.unsqueeze(1)
                 lx = lx * mask
 
                 # scaling for rank dropout: treat as if the rank is changed
@@ -1445,4 +1448,4 @@ class LoRANetwork(torch.nn.Module):
             scalednorm = updown.norm() * ratio
             norms.append(scalednorm.item())
 
-        return keys_scaled, sum(norms) / len(norms), max(norms)
+        return keys_scaled, sum(norms) / len(norms), max(norms)