fixed huge flux training bug. Added ability to use an assistatn lora

toolkit/assistant_lora.py

@@ -0,0 +1,55 @@
+from typing import TYPE_CHECKING
+from toolkit.config_modules import NetworkConfig
+from toolkit.lora_special import LoRASpecialNetwork
+from safetensors.torch import load_file
+
+if TYPE_CHECKING:
+    from toolkit.stable_diffusion_model import StableDiffusion
+
+
+def load_assistant_lora_from_path(adapter_path, sd: 'StableDiffusion') -> LoRASpecialNetwork:
+    if not sd.is_flux:
+        raise ValueError("Only Flux models can load assistant adapters currently.")
+    pipe = sd.pipeline
+    print(f"Loading assistant adapter from {adapter_path}")
+    adapter_name = adapter_path.split("/")[-1].split(".")[0]
+    lora_state_dict = load_file(adapter_path)
+
+    linear_dim = int(lora_state_dict['transformer.single_transformer_blocks.0.attn.to_k.lora_A.weight'].shape[0])
+    # linear_alpha = int(lora_state_dict['lora_transformer_single_transformer_blocks_0_attn_to_k.alpha'].item())
+    linear_alpha = linear_dim
+    transformer_only = 'transformer.proj_out.alpha' not in lora_state_dict
+    # get dim and scale
+    network_config = NetworkConfig(
+        linear=linear_dim,
+        linear_alpha=linear_alpha,
+        transformer_only=transformer_only,
+    )
+
+    network = LoRASpecialNetwork(
+        text_encoder=pipe.text_encoder,
+        unet=pipe.transformer,
+        lora_dim=network_config.linear,
+        multiplier=1.0,
+        alpha=network_config.linear_alpha,
+        train_unet=True,
+        train_text_encoder=False,
+        is_flux=True,
+        network_config=network_config,
+        network_type=network_config.type,
+        transformer_only=network_config.transformer_only,
+        is_assistant_adapter=True
+    )
+    network.apply_to(
+        pipe.text_encoder,
+        pipe.transformer,
+        apply_text_encoder=False,
+        apply_unet=True
+    )
+    network.force_to(sd.device_torch, dtype=sd.torch_dtype)
+    network.eval()
+    network._update_torch_multiplier()
+    network.load_weights(lora_state_dict)
+    network.is_active = True
+
+    return network