lora retrieval

cd0319-ImportError-libX11.so.6

diffsynth/extensions/ImageQualityMetric/open_clip/factory.py

@@ -5,7 +5,7 @@ import pathlib
 import re
 from copy import deepcopy
 from pathlib import Path
-from turtle import forward
+# from turtle import forward
 from typing import Any, Dict, Optional, Tuple, Union
 
 import torch

diffsynth/models/flux_dit.py

@@ -41,6 +41,30 @@ class RoPEEmbedding(torch.nn.Module):
         emb = torch.cat([self.rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3)
         return emb.unsqueeze(1)
 
+class AdaLayerNorm(torch.nn.Module):
+    def __init__(self, dim, single=False, dual=False):
+        super().__init__()
+        self.single = single
+        self.dual = dual
+        self.linear = torch.nn.Linear(dim, dim * [[6, 2][single], 9][dual])
+        self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb, **kwargs):
+        emb = self.linear(torch.nn.functional.silu(emb),**kwargs)
+        if self.single:
+            scale, shift = emb.unsqueeze(1).chunk(2, dim=2)
+            x = self.norm(x) * (1 + scale) + shift
+            return x
+        elif self.dual:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp, shift_msa2, scale_msa2, gate_msa2 = emb.unsqueeze(1).chunk(9, dim=2)
+            norm_x = self.norm(x)
+            x = norm_x * (1 + scale_msa) + shift_msa
+            norm_x2 = norm_x * (1 + scale_msa2) + shift_msa2
+            return x, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_x2, gate_msa2
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.unsqueeze(1).chunk(6, dim=2)
+            x = self.norm(x) * (1 + scale_msa) + shift_msa
+            return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
 
 class FluxJointAttention(torch.nn.Module):
@@ -70,17 +94,17 @@ class FluxJointAttention(torch.nn.Module):
         xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
         return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
 
-    def forward(self, hidden_states_a, hidden_states_b, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+    def forward(self, hidden_states_a, hidden_states_b, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None, **kwargs):
         batch_size = hidden_states_a.shape[0]
 
         # Part A
-        qkv_a = self.a_to_qkv(hidden_states_a)
+        qkv_a = self.a_to_qkv(hidden_states_a,**kwargs)
         qkv_a = qkv_a.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
         q_a, k_a, v_a = qkv_a.chunk(3, dim=1)
         q_a, k_a = self.norm_q_a(q_a), self.norm_k_a(k_a)
 
         # Part B
-        qkv_b = self.b_to_qkv(hidden_states_b)
+        qkv_b = self.b_to_qkv(hidden_states_b,**kwargs)
         qkv_b = qkv_b.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
         q_b, k_b, v_b = qkv_b.chunk(3, dim=1)
         q_b, k_b = self.norm_q_b(q_b), self.norm_k_b(k_b)
@@ -97,13 +121,25 @@ class FluxJointAttention(torch.nn.Module):
         hidden_states_b, hidden_states_a = hidden_states[:, :hidden_states_b.shape[1]], hidden_states[:, hidden_states_b.shape[1]:]
         if ipadapter_kwargs_list is not None:
             hidden_states_a = interact_with_ipadapter(hidden_states_a, q_a, **ipadapter_kwargs_list)
-        hidden_states_a = self.a_to_out(hidden_states_a)
+        hidden_states_a = self.a_to_out(hidden_states_a,**kwargs)
         if self.only_out_a:
             return hidden_states_a
         else:
-            hidden_states_b = self.b_to_out(hidden_states_b)
+            hidden_states_b = self.b_to_out(hidden_states_b,**kwargs)
             return hidden_states_a, hidden_states_b
 
+class AutoSequential(torch.nn.Sequential):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+    def forward(self, input, **kwargs):
+        for module in self:
+            
+            if isinstance(module, torch.nn.Linear):
+                # print("##"*10)
+                input = module(input, **kwargs)
+            else:
+                input = module(input)
+        return input
 
 
 class FluxJointTransformerBlock(torch.nn.Module):
@@ -120,6 +156,11 @@ class FluxJointTransformerBlock(torch.nn.Module):
             torch.nn.GELU(approximate="tanh"),
             torch.nn.Linear(dim*4, dim)
         )
+        # self.ff_a = AutoSequential(
+        #     torch.nn.Linear(dim, dim*4),
+        #     torch.nn.GELU(approximate="tanh"),
+        #     torch.nn.Linear(dim*4, dim)
+        # )
 
         self.norm2_b = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
         self.ff_b = torch.nn.Sequential(
@@ -127,14 +168,18 @@ class FluxJointTransformerBlock(torch.nn.Module):
             torch.nn.GELU(approximate="tanh"),
             torch.nn.Linear(dim*4, dim)
         )
+        # self.ff_b = AutoSequential(
+        #     torch.nn.Linear(dim, dim*4),
+        #     torch.nn.GELU(approximate="tanh"),
+        #     torch.nn.Linear(dim*4, dim)
+        # )
 
-
-    def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
-        norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb)
-        norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb)
+    def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None, **kwargs):
+        norm_hidden_states_a, gate_msa_a, shift_mlp_a, scale_mlp_a, gate_mlp_a = self.norm1_a(hidden_states_a, emb=temb, **kwargs)
+        norm_hidden_states_b, gate_msa_b, shift_mlp_b, scale_mlp_b, gate_mlp_b = self.norm1_b(hidden_states_b, emb=temb, **kwargs)
 
         # Attention
-        attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b, image_rotary_emb, attn_mask, ipadapter_kwargs_list)
+        attn_output_a, attn_output_b = self.attn(norm_hidden_states_a, norm_hidden_states_b, image_rotary_emb, attn_mask, ipadapter_kwargs_list, **kwargs)
 
         # Part A
         hidden_states_a = hidden_states_a + gate_msa_a * attn_output_a
@@ -149,7 +194,6 @@ class FluxJointTransformerBlock(torch.nn.Module):
         return hidden_states_a, hidden_states_b
 
 
-
 class FluxSingleAttention(torch.nn.Module):
     def __init__(self, dim_a, dim_b, num_heads, head_dim):
         super().__init__()
@@ -170,10 +214,10 @@ class FluxSingleAttention(torch.nn.Module):
         return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
 
 
-    def forward(self, hidden_states, image_rotary_emb):
+    def forward(self, hidden_states, image_rotary_emb, **kwargs):
         batch_size = hidden_states.shape[0]
 
-        qkv_a = self.a_to_qkv(hidden_states)
+        qkv_a = self.a_to_qkv(hidden_states,**kwargs)
         qkv_a = qkv_a.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
         q_a, k_a, v = qkv_a.chunk(3, dim=1)
         q_a, k_a = self.norm_q_a(q_a), self.norm_k_a(k_a)
@@ -195,8 +239,8 @@ class AdaLayerNormSingle(torch.nn.Module):
         self.norm = torch.nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
 
-    def forward(self, x, emb):
-        emb = self.linear(self.silu(emb))
+    def forward(self, x, emb, **kwargs):
+        emb = self.linear(self.silu(emb),**kwargs)
         shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
         x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x, gate_msa
@@ -226,7 +270,7 @@ class FluxSingleTransformerBlock(torch.nn.Module):
         return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
 
 
-    def process_attention(self, hidden_states, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+    def process_attention(self, hidden_states, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None, **kwargs):
         batch_size = hidden_states.shape[0]
 
         qkv = hidden_states.view(batch_size, -1, 3 * self.num_heads, self.head_dim).transpose(1, 2)
@@ -243,17 +287,17 @@ class FluxSingleTransformerBlock(torch.nn.Module):
         return hidden_states
 
 
-    def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None):
+    def forward(self, hidden_states_a, hidden_states_b, temb, image_rotary_emb, attn_mask=None, ipadapter_kwargs_list=None, **kwargs):
         residual = hidden_states_a
-        norm_hidden_states, gate = self.norm(hidden_states_a, emb=temb)
-        hidden_states_a = self.to_qkv_mlp(norm_hidden_states)
+        norm_hidden_states, gate = self.norm(hidden_states_a, emb=temb, **kwargs)
+        hidden_states_a = self.to_qkv_mlp(norm_hidden_states, **kwargs)
         attn_output, mlp_hidden_states = hidden_states_a[:, :, :self.dim * 3], hidden_states_a[:, :, self.dim * 3:]
 
-        attn_output = self.process_attention(attn_output, image_rotary_emb, attn_mask, ipadapter_kwargs_list)
+        attn_output = self.process_attention(attn_output, image_rotary_emb, attn_mask, ipadapter_kwargs_list, **kwargs)
         mlp_hidden_states = torch.nn.functional.gelu(mlp_hidden_states, approximate="tanh")
 
         hidden_states_a = torch.cat([attn_output, mlp_hidden_states], dim=2)
-        hidden_states_a = gate.unsqueeze(1) * self.proj_out(hidden_states_a)
+        hidden_states_a = gate.unsqueeze(1) * self.proj_out(hidden_states_a, **kwargs)
         hidden_states_a = residual + hidden_states_a
 
         return hidden_states_a, hidden_states_b
@@ -267,14 +311,13 @@ class AdaLayerNormContinuous(torch.nn.Module):
         self.linear = torch.nn.Linear(dim, dim * 2, bias=True)
         self.norm = torch.nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False)
 
-    def forward(self, x, conditioning):
-        emb = self.linear(self.silu(conditioning))
+    def forward(self, x, conditioning, **kwargs):
+        emb = self.linear(self.silu(conditioning),**kwargs)
         scale, shift = torch.chunk(emb, 2, dim=1)
         x = self.norm(x) * (1 + scale)[:, None] + shift[:, None]
         return x
 
 
-
 class FluxDiT(torch.nn.Module):
     def __init__(self, disable_guidance_embedder=False):
         super().__init__()
@@ -282,6 +325,8 @@ class FluxDiT(torch.nn.Module):
         self.time_embedder = TimestepEmbeddings(256, 3072)
         self.guidance_embedder = None if disable_guidance_embedder else TimestepEmbeddings(256, 3072)
         self.pooled_text_embedder = torch.nn.Sequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072))
+        
+        # self.pooled_text_embedder = AutoSequential(torch.nn.Linear(768, 3072), torch.nn.SiLU(), torch.nn.Linear(3072, 3072))
         self.context_embedder = torch.nn.Linear(4096, 3072)
         self.x_embedder = torch.nn.Linear(64, 3072)
 
@@ -428,12 +473,12 @@ class FluxDiT(torch.nn.Module):
 
         height, width = hidden_states.shape[-2:]
         hidden_states = self.patchify(hidden_states)
-        hidden_states = self.x_embedder(hidden_states)
+        hidden_states = self.x_embedder(hidden_states,**kwargs)
 
         if entity_prompt_emb is not None and entity_masks is not None:
             prompt_emb, image_rotary_emb, attention_mask = self.process_entity_masks(hidden_states, prompt_emb, entity_prompt_emb, entity_masks, text_ids, image_ids)
         else:
-            prompt_emb = self.context_embedder(prompt_emb)
+            prompt_emb = self.context_embedder(prompt_emb, **kwargs)
             image_rotary_emb = self.pos_embedder(torch.cat((text_ids, image_ids), dim=1))
             attention_mask = None
 
@@ -446,26 +491,26 @@ class FluxDiT(torch.nn.Module):
             if self.training and use_gradient_checkpointing:
                 hidden_states, prompt_emb = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(block),
-                    hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask,
+                    hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask, **kwargs,
                     use_reentrant=False,
                 )
             else:
-                hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask)
+                hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask, **kwargs)
 
         hidden_states = torch.cat([prompt_emb, hidden_states], dim=1)
         for block in self.single_blocks:
             if self.training and use_gradient_checkpointing:
                 hidden_states, prompt_emb = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(block),
-                    hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask,
+                    hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask, **kwargs,
                     use_reentrant=False,
                 )
             else:
-                hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask)
+                hidden_states, prompt_emb = block(hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask, **kwargs)
         hidden_states = hidden_states[:, prompt_emb.shape[1]:]
 
-        hidden_states = self.final_norm_out(hidden_states, conditioning)
-        hidden_states = self.final_proj_out(hidden_states)
+        hidden_states = self.final_norm_out(hidden_states, conditioning, **kwargs)
+        hidden_states = self.final_proj_out(hidden_states, **kwargs)
         hidden_states = self.unpatchify(hidden_states, height, width)
 
         return hidden_states
@@ -606,6 +651,10 @@ class FluxDiTStateDictConverter:
         for name, param in state_dict.items():
             if name.endswith(".weight") or name.endswith(".bias"):
                 suffix = ".weight" if name.endswith(".weight") else ".bias"
+                if "lora_B" in name:
+                    suffix = ".lora_B" + suffix
+                if "lora_A" in name:
+                    suffix = ".lora_A" + suffix
                 prefix = name[:-len(suffix)]
                 if prefix in global_rename_dict:
                     state_dict_[global_rename_dict[prefix] + suffix] = param
@@ -630,29 +679,73 @@ class FluxDiTStateDictConverter:
         for name in list(state_dict_.keys()):
             if "single_blocks." in name and ".a_to_q." in name:
                 mlp = state_dict_.get(name.replace(".a_to_q.", ".proj_in_besides_attn."), None)
+                
                 if mlp is None:
-                    mlp = torch.zeros(4 * state_dict_[name].shape[0],
+                    dim = 4
+                    if 'lora_A' in name:
+                        dim = 1
+                    mlp = torch.zeros(dim * state_dict_[name].shape[0],
                                       *state_dict_[name].shape[1:],
                                       dtype=state_dict_[name].dtype)
                 else:
+                    # print('$$'*10)
+                    # mlp_name = name.replace(".a_to_q.", ".proj_in_besides_attn.")
+                    # print(f'mlp name: {mlp_name}')
+                    # print(f'mlp shape: {mlp.shape}')
                     state_dict_.pop(name.replace(".a_to_q.", ".proj_in_besides_attn."))
-                param = torch.concat([
-                    state_dict_.pop(name),
-                    state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")),
-                    state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")),
-                    mlp,
-                ], dim=0)
+                # print(f'mlp shape: {mlp.shape}')
+                if 'lora_A' in name:
+
+                    param = torch.concat([
+                        state_dict_.pop(name),
+                        state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")),
+                        state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")),
+                        mlp,
+                    ], dim=0)
+                elif 'lora_B' in name:
+                    # create zreo matrix
+                    d, r = state_dict_[name].shape
+                    # print('--'*10)
+                    # print(d, r)
+                    param = torch.zeros((3*d+mlp.shape[0], 3*r+mlp.shape[1]), dtype=state_dict_[name].dtype, device=state_dict_[name].device)
+                    param[:d, :r] = state_dict_.pop(name)
+                    param[d:2*d, r:2*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_k."))
+                    param[2*d:3*d, 2*r:3*r] = state_dict_.pop(name.replace(".a_to_q.", ".a_to_v."))
+                    param[3*d:, 3*r:] = mlp
+                else:
+                    param = torch.concat([
+                        state_dict_.pop(name),
+                        state_dict_.pop(name.replace(".a_to_q.", ".a_to_k.")),
+                        state_dict_.pop(name.replace(".a_to_q.", ".a_to_v.")),
+                        mlp,
+                    ], dim=0)
                 name_ = name.replace(".a_to_q.", ".to_qkv_mlp.")
                 state_dict_[name_] = param
         for name in list(state_dict_.keys()):
             for component in ["a", "b"]:
                 if f".{component}_to_q." in name:
                     name_ = name.replace(f".{component}_to_q.", f".{component}_to_qkv.")
-                    param = torch.concat([
-                        state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")],
-                        state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")],
-                        state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")],
-                    ], dim=0)
+                    concat_dim = 0
+                    if 'lora_A' in name:
+                        param = torch.concat([
+                            state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")],
+                            state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")],
+                            state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")],
+                        ], dim=0)
+                    elif 'lora_B' in name:
+                        origin = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")]
+                        d, r = origin.shape
+                        # print(d, r)
+                        param = torch.zeros((3*d, 3*r), dtype=origin.dtype, device=origin.device)
+                        param[:d, :r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")]
+                        param[d:2*d, r:2*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")]
+                        param[2*d:3*d, 2*r:3*r] = state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")]
+                    else:
+                        param = torch.concat([
+                            state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_q.")],
+                            state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_k.")],
+                            state_dict_[name.replace(f".{component}_to_q.", f".{component}_to_v.")],
+                        ], dim=0)
                     state_dict_[name_] = param
                     state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_q."))
                     state_dict_.pop(name.replace(f".{component}_to_q.", f".{component}_to_k."))
@@ -718,22 +811,48 @@ class FluxDiTStateDictConverter:
             "norm.query_norm.scale": "norm_q_a.weight",
         }
         state_dict_ = {}
+
+
         for name, param in state_dict.items():
+            # match lora load
+            l_name = ''
+            if 'lora_A' in name :
+                l_name = 'lora_A'
+            if 'lora_B' in name :
+                l_name = 'lora_B'
+            if l_name != '':
+                name = name.replace(l_name+'.', '')
+            
+
             if name.startswith("model.diffusion_model."):
                 name = name[len("model.diffusion_model."):]
             names = name.split(".")
             if name in rename_dict:
                 rename = rename_dict[name]
                 if name.startswith("final_layer.adaLN_modulation.1."):
-                    param = torch.concat([param[3072:], param[:3072]], dim=0)
-                state_dict_[rename] = param
+                    if l_name == 'lora_A':
+                        param = torch.concat([param[:,3072:], param[:,:3072]], dim=1)
+                    else:
+                        param = torch.concat([param[3072:], param[:3072]], dim=0)
+                if l_name != '':
+                    state_dict_[rename.replace('weight',l_name+'.weight')] = param
+                else:
+                    state_dict_[rename] = param
+               
             elif names[0] == "double_blocks":
                 rename = f"blocks.{names[1]}." + suffix_rename_dict[".".join(names[2:])]
-                state_dict_[rename] = param
+                if l_name != '':
+                    state_dict_[rename.replace('weight',l_name+'.weight')] = param
+                else:
+                    state_dict_[rename] = param
+
             elif names[0] == "single_blocks":
                 if ".".join(names[2:]) in suffix_rename_dict:
                     rename = f"single_blocks.{names[1]}." + suffix_rename_dict[".".join(names[2:])]
-                    state_dict_[rename] = param
+                    if l_name != '':
+                        state_dict_[rename.replace('weight',l_name+'.weight')] = param
+                    else:
+                        state_dict_[rename] = param
             else:
                 pass
         if "guidance_embedder.timestep_embedder.0.weight" not in state_dict_:

diffsynth/models/lora.py

@@ -26,6 +26,12 @@ class LoRAFromCivitai:
                 return self.convert_state_dict_up_down(state_dict, lora_prefix, alpha)
         return self.convert_state_dict_AB(state_dict, lora_prefix, alpha)
 
+    def convert_state_name(self, state_dict, lora_prefix="lora_unet_", alpha=1.0):
+        for key in state_dict:
+            if ".lora_up" in key:
+                return self.convert_state_name_up_down(state_dict, lora_prefix, alpha)
+        return self.convert_state_name_AB(state_dict, lora_prefix, alpha)
+
 
     def convert_state_dict_up_down(self, state_dict, lora_prefix="lora_unet_", alpha=1.0):
         renamed_lora_prefix = self.renamed_lora_prefix.get(lora_prefix, "")
@@ -50,13 +56,37 @@ class LoRAFromCivitai:
         return state_dict_
     
 
+    def convert_state_name_up_down(self, state_dict, lora_prefix="lora_unet_", alpha=1.0):
+        renamed_lora_prefix = self.renamed_lora_prefix.get(lora_prefix, "")
+        state_dict_ = {}
+        for key in state_dict:
+            if ".lora_up" not in key:
+                continue
+            if not key.startswith(lora_prefix):
+                continue
+            weight_up = state_dict[key].to(device="cuda", dtype=torch.float16)
+            weight_down = state_dict[key.replace(".lora_up", ".lora_down")].to(device="cuda", dtype=torch.float16)
+            if len(weight_up.shape) == 4:
+                weight_up = weight_up.squeeze(3).squeeze(2).to(torch.float32)
+                weight_down = weight_down.squeeze(3).squeeze(2).to(torch.float32)
+            target_name = key.split(".")[0].replace(lora_prefix, renamed_lora_prefix).replace("_", ".") + ".weight"
+            for special_key in self.special_keys:
+                target_name = target_name.replace(special_key, self.special_keys[special_key])
+            
+            state_dict_[target_name.replace(".weight",".lora_B.weight")] = weight_up.cpu()
+            state_dict_[target_name.replace(".weight",".lora_A.weight")] = weight_down.cpu()
+        return state_dict_
+    
+
     def convert_state_dict_AB(self, state_dict, lora_prefix="", alpha=1.0, device="cuda", torch_dtype=torch.float16):
         state_dict_ = {}
+
         for key in state_dict:
             if ".lora_B." not in key:
                 continue
             if not key.startswith(lora_prefix):
                 continue
+
             weight_up = state_dict[key].to(device=device, dtype=torch_dtype)
             weight_down = state_dict[key.replace(".lora_B.", ".lora_A.")].to(device=device, dtype=torch_dtype)
             if len(weight_up.shape) == 4:
@@ -67,11 +97,39 @@ class LoRAFromCivitai:
                 lora_weight = alpha * torch.mm(weight_up, weight_down)
             keys = key.split(".")
             keys.pop(keys.index("lora_B"))
+            
             target_name = ".".join(keys)
+           
             target_name = target_name[len(lora_prefix):]
+           
             state_dict_[target_name] = lora_weight.cpu()
         return state_dict_
     
+    def convert_state_name_AB(self, state_dict, lora_prefix="", alpha=1.0, device="cuda", torch_dtype=torch.float16):
+        state_dict_ = {}
+ 
+        for key in state_dict:
+            if ".lora_B." not in key:
+                continue
+            if not key.startswith(lora_prefix):
+                continue
+
+            weight_up = state_dict[key].to(device=device, dtype=torch_dtype)
+            weight_down = state_dict[key.replace(".lora_B.", ".lora_A.")].to(device=device, dtype=torch_dtype)
+            if len(weight_up.shape) == 4:
+                weight_up = weight_up.squeeze(3).squeeze(2)
+                weight_down = weight_down.squeeze(3).squeeze(2)
+
+            keys = key.split(".")
+            keys.pop(keys.index("lora_B"))
+
+            target_name = ".".join(keys)
+
+            target_name = target_name[len(lora_prefix):]
+
+            state_dict_[target_name.replace(".weight",".lora_B.weight")] = weight_up.cpu()
+            state_dict_[target_name.replace(".weight",".lora_A.weight")] = weight_down.cpu()
+        return state_dict_
 
     def load(self, model, state_dict_lora, lora_prefix, alpha=1.0, model_resource=None):
         state_dict_model = model.state_dict()
@@ -100,13 +158,16 @@ class LoRAFromCivitai:
         for lora_prefix, model_class in zip(self.lora_prefix, self.supported_model_classes):
             if not isinstance(model, model_class):
                 continue
+            # print(f'lora_prefix: {lora_prefix}')
             state_dict_model = model.state_dict()
             for model_resource in ["diffusers", "civitai"]:
                 try:
                     state_dict_lora_ = self.convert_state_dict(state_dict_lora, lora_prefix=lora_prefix, alpha=1.0)
+                    # print(f'after convert_state_dict lora state_dict:{state_dict_lora_.keys()}')
                     converter_fn = model.__class__.state_dict_converter().from_diffusers if model_resource == "diffusers" \
                         else model.__class__.state_dict_converter().from_civitai
                     state_dict_lora_ = converter_fn(state_dict_lora_)
+                    # print(f'after converter_fn lora state_dict:{state_dict_lora_.keys()}')
                     if isinstance(state_dict_lora_, tuple):
                         state_dict_lora_ = state_dict_lora_[0]
                     if len(state_dict_lora_) == 0:
@@ -120,7 +181,35 @@ class LoRAFromCivitai:
                     pass
         return None
 
+    def get_converted_lora_state_dict(self, model, state_dict_lora):
+        for lora_prefix, model_class in zip(self.lora_prefix, self.supported_model_classes):
+            if not isinstance(model, model_class):
+                continue
 
+            state_dict_model = model.state_dict()
+            for model_resource in ["diffusers","civitai"]:
+                try:
+                    state_dict_lora_ = self.convert_state_name(state_dict_lora, lora_prefix=lora_prefix, alpha=1.0)
+ 
+                    converter_fn = model.__class__.state_dict_converter().from_diffusers if model_resource == 'diffusers' \
+                        else model.__class__.state_dict_converter().from_civitai
+                    state_dict_lora_ = converter_fn(state_dict_lora_)
+                    
+                    if isinstance(state_dict_lora_, tuple):
+                        state_dict_lora_ = state_dict_lora_[0]
+
+                    if len(state_dict_lora_) == 0:
+                        continue
+                    # return state_dict_lora_
+                    for name in state_dict_lora_:
+                        if name.replace('.lora_B','').replace('.lora_A','') not in state_dict_model:
+                            print(f"   lora's {name} is not in model.")
+                            break
+                    else:
+                        return state_dict_lora_
+                except Exception as e:
+                    print(f"error   {str(e)}")
+        return None
 
 class SDLoRAFromCivitai(LoRAFromCivitai):
     def __init__(self):
@@ -195,73 +284,85 @@ class FluxLoRAFromCivitai(LoRAFromCivitai):
             "txt.mod": "txt_mod",
         }
 
-    
-    
+
 class GeneralLoRAFromPeft:
     def __init__(self):
         self.supported_model_classes = [SDUNet, SDXLUNet, SD3DiT, HunyuanDiT, FluxDiT, CogDiT, WanModel]
-    
-    
-    def get_name_dict(self, lora_state_dict):
-        lora_name_dict = {}
-        for key in lora_state_dict:
+
+
+    def fetch_device_dtype_from_state_dict(self, state_dict):
+        device, torch_dtype = None, None
+        for name, param in state_dict.items():
+            device, torch_dtype = param.device, param.dtype
+            break
+        return device, torch_dtype
+
+
+    def convert_state_dict(self, state_dict, alpha=1.0, target_state_dict={}):
+        device, torch_dtype = self.fetch_device_dtype_from_state_dict(target_state_dict)
+        if torch_dtype == torch.float8_e4m3fn:
+            torch_dtype = torch.float32
+        state_dict_ = {}
+        for key in state_dict:
             if ".lora_B." not in key:
                 continue
+            weight_up = state_dict[key].to(device=device, dtype=torch_dtype)
+            weight_down = state_dict[key.replace(".lora_B.", ".lora_A.")].to(device=device, dtype=torch_dtype)
+            if len(weight_up.shape) == 4:
+                weight_up = weight_up.squeeze(3).squeeze(2)
+                weight_down = weight_down.squeeze(3).squeeze(2)
+                lora_weight = alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
+            else:
+                lora_weight = alpha * torch.mm(weight_up, weight_down)
             keys = key.split(".")
             if len(keys) > keys.index("lora_B") + 2:
                 keys.pop(keys.index("lora_B") + 1)
             keys.pop(keys.index("lora_B"))
-            if keys[0] == "diffusion_model":
-                keys.pop(0)
             target_name = ".".join(keys)
-            lora_name_dict[target_name] = (key, key.replace(".lora_B.", ".lora_A."))
-        return lora_name_dict
+            if target_name.startswith("diffusion_model."):
+                target_name = target_name[len("diffusion_model."):]
+            if target_name not in target_state_dict:
+                return {}
+            state_dict_[target_name] = lora_weight.cpu()
+        return state_dict_
     
-    
-    def match(self, model: torch.nn.Module, state_dict_lora):
-        lora_name_dict = self.get_name_dict(state_dict_lora)
-        model_name_dict = {name: None for name, _ in model.named_parameters()}
-        matched_num = sum([i in model_name_dict for i in lora_name_dict])
-        if matched_num == len(lora_name_dict):
-            return "", ""
-        else:
-            return None
-    
-    
-    def fetch_device_and_dtype(self, state_dict):
-        device, dtype = None, None
-        for name, param in state_dict.items():
-            device, dtype = param.device, param.dtype
-            break
-        computation_device = device
-        computation_dtype = dtype
-        if computation_device == torch.device("cpu"):
-            if torch.cuda.is_available():
-                computation_device = torch.device("cuda")
-        if computation_dtype == torch.float8_e4m3fn:
-            computation_dtype = torch.float32
-        return device, dtype, computation_device, computation_dtype
-
 
     def load(self, model, state_dict_lora, lora_prefix="", alpha=1.0, model_resource=""):
         state_dict_model = model.state_dict()
-        device, dtype, computation_device, computation_dtype = self.fetch_device_and_dtype(state_dict_model)
-        lora_name_dict = self.get_name_dict(state_dict_lora)
-        for name in lora_name_dict:
-            weight_up = state_dict_lora[lora_name_dict[name][0]].to(device=computation_device, dtype=computation_dtype)
-            weight_down = state_dict_lora[lora_name_dict[name][1]].to(device=computation_device, dtype=computation_dtype)
-            if len(weight_up.shape) == 4:
-                weight_up = weight_up.squeeze(3).squeeze(2)
-                weight_down = weight_down.squeeze(3).squeeze(2)
-                weight_lora = alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3)
-            else:
-                weight_lora = alpha * torch.mm(weight_up, weight_down)
-            weight_model = state_dict_model[name].to(device=computation_device, dtype=computation_dtype)
-            weight_patched = weight_model + weight_lora
-            state_dict_model[name] = weight_patched.to(device=device, dtype=dtype)
-        print(f"    {len(lora_name_dict)} tensors are updated.")
-        model.load_state_dict(state_dict_model)
+        state_dict_lora = self.convert_state_dict(state_dict_lora, alpha=alpha, target_state_dict=state_dict_model)
+        if len(state_dict_lora) > 0:
+            print(f"    {len(state_dict_lora)} tensors are updated.")
+            for name in state_dict_lora:
+                if state_dict_model[name].dtype == torch.float8_e4m3fn:
+                    weight = state_dict_model[name].to(torch.float32)
+                    lora_weight = state_dict_lora[name].to(
+                        dtype=torch.float32,
+                        device=state_dict_model[name].device
+                    )
+                    state_dict_model[name] = (weight + lora_weight).to(
+                        dtype=state_dict_model[name].dtype,
+                        device=state_dict_model[name].device
+                    )
+                else:
+                    state_dict_model[name] += state_dict_lora[name].to(
+                        dtype=state_dict_model[name].dtype,
+                        device=state_dict_model[name].device
+                    )
+            model.load_state_dict(state_dict_model)
     
+
+    def match(self, model, state_dict_lora):
+        for model_class in self.supported_model_classes:
+            if not isinstance(model, model_class):
+                continue
+            state_dict_model = model.state_dict()
+            try:
+                state_dict_lora_ = self.convert_state_dict(state_dict_lora, alpha=1.0, target_state_dict=state_dict_model)
+                if len(state_dict_lora_) > 0:
+                    return "", ""
+            except:
+                pass
+        return None
     
 
 class HunyuanVideoLoRAFromCivitai(LoRAFromCivitai):

diffsynth/models/utils.py

@@ -62,25 +62,26 @@ def load_state_dict_from_folder(file_path, torch_dtype=None):
     return state_dict
 
 
-def load_state_dict(file_path, torch_dtype=None):
+def load_state_dict(file_path, torch_dtype=None, device="cpu"):
     if file_path.endswith(".safetensors"):
-        return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype)
+        return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device)
     else:
-        return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype)
+        return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device)
 
 
-def load_state_dict_from_safetensors(file_path, torch_dtype=None):
+def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"):
     state_dict = {}
     with safe_open(file_path, framework="pt", device="cpu") as f:
         for k in f.keys():
             state_dict[k] = f.get_tensor(k)
             if torch_dtype is not None:
                 state_dict[k] = state_dict[k].to(torch_dtype)
+            state_dict[k] = state_dict[k].to(device)
     return state_dict
 
 
-def load_state_dict_from_bin(file_path, torch_dtype=None):
-    state_dict = torch.load(file_path, map_location="cpu", weights_only=True)
+def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"):
+    state_dict = torch.load(file_path, map_location=device, weights_only=True)
     if torch_dtype is not None:
         for i in state_dict:
             if isinstance(state_dict[i], torch.Tensor):

diffsynth/models/wan_video_controlnet.py

@@ -1,204 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import Tuple, Optional
-from einops import rearrange
-from .wan_video_dit import DiTBlock, precompute_freqs_cis_3d, MLP, sinusoidal_embedding_1d
-from .utils import hash_state_dict_keys
-
-
-
-class WanControlNetModel(torch.nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        in_dim: int,
-        ffn_dim: int,
-        out_dim: int,
-        text_dim: int,
-        freq_dim: int,
-        eps: float,
-        patch_size: Tuple[int, int, int],
-        num_heads: int,
-        num_layers: int,
-        has_image_input: bool,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.freq_dim = freq_dim
-        self.has_image_input = has_image_input
-        self.patch_size = patch_size
-
-        self.patch_embedding = nn.Conv3d(
-            in_dim, dim, kernel_size=patch_size, stride=patch_size)
-        self.text_embedding = nn.Sequential(
-            nn.Linear(text_dim, dim),
-            nn.GELU(approximate='tanh'),
-            nn.Linear(dim, dim)
-        )
-        self.time_embedding = nn.Sequential(
-            nn.Linear(freq_dim, dim),
-            nn.SiLU(),
-            nn.Linear(dim, dim)
-        )
-        self.time_projection = nn.Sequential(
-            nn.SiLU(), nn.Linear(dim, dim * 6))
-        self.blocks = nn.ModuleList([
-            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps)
-            for _ in range(num_layers)
-        ])
-        head_dim = dim // num_heads
-        self.freqs = precompute_freqs_cis_3d(head_dim)
-
-        if has_image_input:
-            self.img_emb = MLP(1280, dim)  # clip_feature_dim = 1280
-            
-        self.controlnet_conv_in = torch.nn.Conv3d(in_channels=in_dim, out_channels=in_dim, kernel_size=1)
-        self.controlnet_blocks = torch.nn.ModuleList([
-            torch.nn.Linear(dim, dim, bias=False)
-            for _ in range(num_layers)
-        ])
-
-    def patchify(self, x: torch.Tensor):
-        x = self.patch_embedding(x)
-        grid_size = x.shape[2:]
-        x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous()
-        return x, grid_size  # x, grid_size: (f, h, w)
-
-    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor):
-        return rearrange(
-            x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)',
-            f=grid_size[0], h=grid_size[1], w=grid_size[2], 
-            x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2]
-        )
-
-    def forward(self,
-                x: torch.Tensor,
-                timestep: torch.Tensor,
-                context: torch.Tensor,
-                clip_feature: Optional[torch.Tensor] = None,
-                y: Optional[torch.Tensor] = None,
-                controlnet_conditioning: Optional[torch.Tensor] = None,
-                use_gradient_checkpointing: bool = False,
-                use_gradient_checkpointing_offload: bool = False,
-                **kwargs,
-                ):
-        t = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, timestep))
-        t_mod = self.time_projection(t).unflatten(1, (6, self.dim))
-        context = self.text_embedding(context)
-        
-        if self.has_image_input:
-            x = torch.cat([x, y], dim=1)  # (b, c_x + c_y, f, h, w)
-            clip_embdding = self.img_emb(clip_feature)
-            context = torch.cat([clip_embdding, context], dim=1)
-            
-        x = x + self.controlnet_conv_in(controlnet_conditioning)
-        x, (f, h, w) = self.patchify(x)
-        
-        freqs = torch.cat([
-            self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
-            self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
-            self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
-        ], dim=-1).reshape(f * h * w, 1, -1).to(x.device)
-        
-        def create_custom_forward(module):
-            def custom_forward(*inputs):
-                return module(*inputs)
-            return custom_forward
-
-        res_stack = []
-        for block in self.blocks:
-            if self.training and use_gradient_checkpointing:
-                if use_gradient_checkpointing_offload:
-                    with torch.autograd.graph.save_on_cpu():
-                        x = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(block),
-                            x, context, t_mod, freqs,
-                            use_reentrant=False,
-                        )
-                else:
-                    x = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(block),
-                        x, context, t_mod, freqs,
-                        use_reentrant=False,
-                    )
-            else:
-                x = block(x, context, t_mod, freqs)
-            res_stack.append(x)
-        
-        controlnet_res_stack = [block(res) for block, res in zip(self.controlnet_blocks, res_stack)]
-        return controlnet_res_stack
-
-    @staticmethod
-    def state_dict_converter():
-        return WanControlNetModelStateDictConverter()
-    
-    
-class WanControlNetModelStateDictConverter:
-    def __init__(self):
-        pass
-    
-    def from_diffusers(self, state_dict):
-        return state_dict
-    
-    def from_civitai(self, state_dict):
-        return state_dict
-    
-    def from_base_model(self, state_dict):
-        if hash_state_dict_keys(state_dict) == "9269f8db9040a9d860eaca435be61814":
-            config = {
-                "has_image_input": False,
-                "patch_size": [1, 2, 2],
-                "in_dim": 16,
-                "dim": 1536,
-                "ffn_dim": 8960,
-                "freq_dim": 256,
-                "text_dim": 4096,
-                "out_dim": 16,
-                "num_heads": 12,
-                "num_layers": 30,
-                "eps": 1e-6
-            }
-        elif hash_state_dict_keys(state_dict) == "aafcfd9672c3a2456dc46e1cb6e52c70":
-            config = {
-                "has_image_input": False,
-                "patch_size": [1, 2, 2],
-                "in_dim": 16,
-                "dim": 5120,
-                "ffn_dim": 13824,
-                "freq_dim": 256,
-                "text_dim": 4096,
-                "out_dim": 16,
-                "num_heads": 40,
-                "num_layers": 40,
-                "eps": 1e-6
-            }
-        elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e":
-            config = {
-                "has_image_input": True,
-                "patch_size": [1, 2, 2],
-                "in_dim": 36,
-                "dim": 5120,
-                "ffn_dim": 13824,
-                "freq_dim": 256,
-                "text_dim": 4096,
-                "out_dim": 16,
-                "num_heads": 40,
-                "num_layers": 40,
-                "eps": 1e-6
-            }
-        else:
-            config = {}
-        state_dict_ = {}
-        dtype, device = None, None
-        for name, param in state_dict.items():
-            if name.startswith("head."):
-                continue
-            state_dict_[name] = param
-            dtype, device = param.dtype, param.device
-        for block_id in range(config["num_layers"]):
-            zeros = torch.zeros((config["dim"], config["dim"]), dtype=dtype, device=device)
-            state_dict_[f"controlnet_blocks.{block_id}.weight"] = zeros.clone()
-        state_dict_["controlnet_conv_in.weight"] = torch.zeros((config["in_dim"], config["in_dim"], 1, 1, 1), dtype=dtype, device=device)
-        state_dict_["controlnet_conv_in.bias"] = torch.zeros((config["in_dim"],), dtype=dtype, device=device)
-        return state_dict_, config

diffsynth/models/wan_video_motion_controller.py

@@ -1,27 +0,0 @@
-import torch
-import torch.nn as nn
-from .wan_video_dit import sinusoidal_embedding_1d
-
-
-
-class WanMotionControllerModel(torch.nn.Module):
-    def __init__(self, freq_dim=256, dim=1536):
-        super().__init__()
-        self.freq_dim = freq_dim
-        self.linear = nn.Sequential(
-            nn.Linear(freq_dim, dim),
-            nn.SiLU(),
-            nn.Linear(dim, dim),
-            nn.SiLU(),
-            nn.Linear(dim, dim * 6),
-        )
-        
-    def forward(self, motion_bucket_id):
-        emb = sinusoidal_embedding_1d(self.freq_dim, motion_bucket_id * 10)
-        emb = self.linear(emb)
-        return emb
-    
-    def init(self):
-        state_dict = self.linear[-1].state_dict()
-        state_dict = {i: state_dict[i] * 0 for i in state_dict}
-        self.linear[-1].load_state_dict(state_dict)

diffsynth/pipelines/flux_image.py

@@ -13,7 +13,7 @@ from transformers import SiglipVisionModel
 from copy import deepcopy
 from transformers.models.t5.modeling_t5 import T5LayerNorm, T5DenseActDense, T5DenseGatedActDense
 from ..models.flux_dit import RMSNorm
-from ..vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear
+from ..vram_management import enable_vram_management, enable_auto_lora, AutoLoRALinear, AutoWrappedModule, AutoWrappedLinear
 
 
 class FluxImagePipeline(BasePipeline):
@@ -132,6 +132,15 @@ class FluxImagePipeline(BasePipeline):
         )
         self.enable_cpu_offload()
 
+    def enable_auto_lora(self):
+        enable_auto_lora(
+            self.dit,
+            module_map={
+                RMSNorm: AutoWrappedModule,
+                torch.nn.Linear: AutoLoRALinear,
+            },
+            name_prefix=''
+        )
 
     def denoising_model(self):
         return self.dit
@@ -391,6 +400,9 @@ class FluxImagePipeline(BasePipeline):
         # Progress bar
         progress_bar_cmd=tqdm,
         progress_bar_st=None,
+        lora_state_dicts=[],
+        lora_alphas=[],
+        lora_patcher=None,
     ):
         height, width = self.check_resize_height_width(height, width)
 
@@ -430,6 +442,9 @@ class FluxImagePipeline(BasePipeline):
             inference_callback = lambda prompt_emb_posi, controlnet_kwargs: lets_dance_flux(
                 dit=self.dit, controlnet=self.controlnet,
                 hidden_states=latents, timestep=timestep,
+                lora_state_dicts=lora_state_dicts,
+                lora_alphas = lora_alphas,
+                lora_patcher=lora_patcher,
                 **prompt_emb_posi, **tiler_kwargs, **extra_input, **controlnet_kwargs, **ipadapter_kwargs_list_posi, **eligen_kwargs_posi, **tea_cache_kwargs,
             )
             noise_pred_posi = self.control_noise_via_local_prompts(
@@ -447,6 +462,9 @@ class FluxImagePipeline(BasePipeline):
                 noise_pred_nega = lets_dance_flux(
                     dit=self.dit, controlnet=self.controlnet,
                     hidden_states=latents, timestep=timestep,
+                    lora_state_dicts=lora_state_dicts,
+                    lora_alphas = lora_alphas,
+                    lora_patcher=lora_patcher,
                     **prompt_emb_nega, **tiler_kwargs, **extra_input, **controlnet_kwargs_nega, **ipadapter_kwargs_list_nega, **eligen_kwargs_nega,
                 )
                 noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega)
@@ -511,7 +529,6 @@ class TeaCache:
         hidden_states = hidden_states + self.previous_residual
         return hidden_states
 
-
 def lets_dance_flux(
     dit: FluxDiT,
     controlnet: FluxMultiControlNetManager = None,
@@ -530,8 +547,10 @@ def lets_dance_flux(
     entity_masks=None,
     ipadapter_kwargs_list={},
     tea_cache: TeaCache = None,
+    use_gradient_checkpointing=False,
     **kwargs
 ):
+
     if tiled:
         def flux_forward_fn(hl, hr, wl, wr):
             tiled_controlnet_frames = [f[:, :, hl: hr, wl: wr] for f in controlnet_frames] if controlnet_frames is not None else None
@@ -595,6 +614,11 @@ def lets_dance_flux(
         prompt_emb = dit.context_embedder(prompt_emb)
         image_rotary_emb = dit.pos_embedder(torch.cat((text_ids, image_ids), dim=1))
         attention_mask = None
+        
+    def create_custom_forward(module):
+        def custom_forward(*inputs, **kwargs):
+            return module(*inputs, **kwargs)
+        return custom_forward
 
     # TeaCache
     if tea_cache is not None:
@@ -607,14 +631,22 @@ def lets_dance_flux(
     else:
         # Joint Blocks
         for block_id, block in enumerate(dit.blocks):
-            hidden_states, prompt_emb = block(
-                hidden_states,
-                prompt_emb,
-                conditioning,
-                image_rotary_emb,
-                attention_mask,
-                ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id, None)
-            )
+            if use_gradient_checkpointing:
+                hidden_states, prompt_emb = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask, ipadapter_kwargs_list.get(block_id, None), **kwargs,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, prompt_emb = block(
+                    hidden_states,
+                    prompt_emb,
+                    conditioning,
+                    image_rotary_emb,
+                    attention_mask,
+                    ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id, None),
+                    **kwargs
+                )
             # ControlNet
             if controlnet is not None and controlnet_frames is not None:
                 hidden_states = hidden_states + controlnet_res_stack[block_id]
@@ -623,14 +655,22 @@ def lets_dance_flux(
         hidden_states = torch.cat([prompt_emb, hidden_states], dim=1)
         num_joint_blocks = len(dit.blocks)
         for block_id, block in enumerate(dit.single_blocks):
-            hidden_states, prompt_emb = block(
-                hidden_states,
-                prompt_emb,
-                conditioning,
-                image_rotary_emb,
-                attention_mask,
-                ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id + num_joint_blocks, None)
-            )
+            if use_gradient_checkpointing:
+                hidden_states, prompt_emb = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states, prompt_emb, conditioning, image_rotary_emb, attention_mask, ipadapter_kwargs_list.get(block_id + num_joint_blocks, None), **kwargs,
+                    use_reentrant=False,
+                )
+            else:
+                hidden_states, prompt_emb = block(
+                    hidden_states,
+                    prompt_emb,
+                    conditioning,
+                    image_rotary_emb,
+                    attention_mask,
+                    ipadapter_kwargs_list=ipadapter_kwargs_list.get(block_id + num_joint_blocks, None),
+                    **kwargs
+                )
             # ControlNet
             if controlnet is not None and controlnet_frames is not None:
                 hidden_states[:, prompt_emb.shape[1]:] = hidden_states[:, prompt_emb.shape[1]:] + controlnet_single_res_stack[block_id]
@@ -639,8 +679,8 @@ def lets_dance_flux(
         if tea_cache is not None:
             tea_cache.store(hidden_states)
 
-    hidden_states = dit.final_norm_out(hidden_states, conditioning)
-    hidden_states = dit.final_proj_out(hidden_states)
+    hidden_states = dit.final_norm_out(hidden_states, conditioning, **kwargs)
+    hidden_states = dit.final_proj_out(hidden_states, **kwargs)
     hidden_states = dit.unpatchify(hidden_states, height, width)
 
     return hidden_states

diffsynth/pipelines/wan_video.py

@@ -17,8 +17,6 @@ from ..vram_management import enable_vram_management, AutoWrappedModule, AutoWra
 from ..models.wan_video_text_encoder import T5RelativeEmbedding, T5LayerNorm
 from ..models.wan_video_dit import RMSNorm, sinusoidal_embedding_1d
 from ..models.wan_video_vae import RMS_norm, CausalConv3d, Upsample
-from ..models.wan_video_controlnet import WanControlNetModel
-from ..models.wan_video_motion_controller import WanMotionControllerModel
 
 
 
@@ -32,9 +30,7 @@ class WanVideoPipeline(BasePipeline):
         self.image_encoder: WanImageEncoder = None
         self.dit: WanModel = None
         self.vae: WanVideoVAE = None
-        self.controlnet: WanControlNetModel = None
-        self.motion_controller: WanMotionControllerModel = None
-        self.model_names = ['text_encoder', 'dit', 'vae', 'image_encoder', 'controlnet', 'motion_controller']
+        self.model_names = ['text_encoder', 'dit', 'vae']
         self.height_division_factor = 16
         self.width_division_factor = 16
 
@@ -193,16 +189,6 @@ class WanVideoPipeline(BasePipeline):
     def decode_video(self, latents, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)):
         frames = self.vae.decode(latents, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
         return frames
-    
-    
-    def prepare_controlnet(self, controlnet_frames, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)):
-        controlnet_conditioning = self.encode_video(controlnet_frames, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride).to(dtype=self.torch_dtype, device=self.device)
-        return {"controlnet_conditioning": controlnet_conditioning}
-    
-    
-    def prepare_motion_bucket_id(self, motion_bucket_id):
-        motion_bucket_id = torch.Tensor((motion_bucket_id,)).to(dtype=self.torch_dtype, device=self.device)
-        return {"motion_bucket_id": motion_bucket_id}
 
 
     @torch.no_grad()
@@ -221,13 +207,11 @@ class WanVideoPipeline(BasePipeline):
         cfg_scale=5.0,
         num_inference_steps=50,
         sigma_shift=5.0,
-        motion_bucket_id=None,
         tiled=True,
         tile_size=(30, 52),
         tile_stride=(15, 26),
         tea_cache_l1_thresh=None,
         tea_cache_model_id="",
-        controlnet_frames=None,
         progress_bar_cmd=tqdm,
         progress_bar_st=None,
     ):
@@ -268,21 +252,6 @@ class WanVideoPipeline(BasePipeline):
         else:
             image_emb = {}
             
-        # ControlNet
-        if self.controlnet is not None and controlnet_frames is not None:
-            self.load_models_to_device(['vae', 'controlnet'])
-            controlnet_frames = self.preprocess_images(controlnet_frames)
-            controlnet_frames = torch.stack(controlnet_frames, dim=2).to(dtype=self.torch_dtype, device=self.device)
-            controlnet_kwargs = self.prepare_controlnet(controlnet_frames)
-        else:
-            controlnet_kwargs = {}
-            
-        # Motion Controller
-        if self.motion_controller is not None and motion_bucket_id is not None:
-            motion_kwargs = self.prepare_motion_bucket_id(motion_bucket_id)
-        else:
-            motion_kwargs = {}
-            
         # Extra input
         extra_input = self.prepare_extra_input(latents)
         
@@ -291,24 +260,14 @@ class WanVideoPipeline(BasePipeline):
         tea_cache_nega = {"tea_cache": TeaCache(num_inference_steps, rel_l1_thresh=tea_cache_l1_thresh, model_id=tea_cache_model_id) if tea_cache_l1_thresh is not None else None}
 
         # Denoise
-        self.load_models_to_device(["dit", "controlnet", "motion_controller"])
+        self.load_models_to_device(["dit"])
         for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):
             timestep = timestep.unsqueeze(0).to(dtype=self.torch_dtype, device=self.device)
 
             # Inference
-            noise_pred_posi = model_fn_wan_video(
-                self.dit, controlnet=self.controlnet, motion_controller=self.motion_controller,
-                x=latents, timestep=timestep,
-                **prompt_emb_posi, **image_emb, **extra_input,
-                **tea_cache_posi, **controlnet_kwargs, **motion_kwargs,
-            )
+            noise_pred_posi = model_fn_wan_video(self.dit, latents, timestep=timestep, **prompt_emb_posi, **image_emb, **extra_input, **tea_cache_posi)
             if cfg_scale != 1.0:
-                noise_pred_nega = model_fn_wan_video(
-                    self.dit, controlnet=self.controlnet, motion_controller=self.motion_controller,
-                    x=latents, timestep=timestep,
-                    **prompt_emb_nega, **image_emb, **extra_input,
-                    **tea_cache_nega, **controlnet_kwargs, **motion_kwargs,
-                )
+                noise_pred_nega = model_fn_wan_video(self.dit, latents, timestep=timestep, **prompt_emb_nega, **image_emb, **extra_input, **tea_cache_nega)
                 noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega)
             else:
                 noise_pred = noise_pred_posi
@@ -381,35 +340,16 @@ class TeaCache:
 
 def model_fn_wan_video(
     dit: WanModel,
-    controlnet: WanControlNetModel = None,
-    motion_controller: WanMotionControllerModel = None,
-    x: torch.Tensor = None,
-    timestep: torch.Tensor = None,
-    context: torch.Tensor = None,
+    x: torch.Tensor,
+    timestep: torch.Tensor,
+    context: torch.Tensor,
     clip_feature: Optional[torch.Tensor] = None,
     y: Optional[torch.Tensor] = None,
     tea_cache: TeaCache = None,
-    controlnet_conditioning: Optional[torch.Tensor] = None,
-    motion_bucket_id: Optional[torch.Tensor] = None,
-    use_gradient_checkpointing: bool = False,
-    use_gradient_checkpointing_offload: bool = False,
     **kwargs,
 ):
-    # ControlNet
-    if controlnet is not None and controlnet_conditioning is not None:
-        controlnet_res_stack = controlnet(
-            x, timestep=timestep, context=context, clip_feature=clip_feature, y=y,
-            controlnet_conditioning=controlnet_conditioning,
-            use_gradient_checkpointing=use_gradient_checkpointing,
-            use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
-        )
-    else:
-        controlnet_res_stack = None
-    
     t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep))
     t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim))
-    if motion_bucket_id is not None and motion_controller is not None:
-        t_mod = t_mod + motion_controller(motion_bucket_id).unflatten(1, (6, dit.dim))
     context = dit.text_embedding(context)
     
     if dit.has_image_input:
@@ -430,35 +370,13 @@ def model_fn_wan_video(
         tea_cache_update = tea_cache.check(dit, x, t_mod)
     else:
         tea_cache_update = False
-        
-    def create_custom_forward(module):
-        def custom_forward(*inputs):
-            return module(*inputs)
-        return custom_forward
     
     if tea_cache_update:
         x = tea_cache.update(x)
     else:
         # blocks
-        for block_id, block in enumerate(dit.blocks):
-            if dit.training and use_gradient_checkpointing:
-                if use_gradient_checkpointing_offload:
-                    with torch.autograd.graph.save_on_cpu():
-                        x = torch.utils.checkpoint.checkpoint(
-                            create_custom_forward(block),
-                            x, context, t_mod, freqs,
-                            use_reentrant=False,
-                        )
-                else:
-                    x = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(block),
-                        x, context, t_mod, freqs,
-                        use_reentrant=False,
-                    )
-            else:
-                x = block(x, context, t_mod, freqs)
-            if controlnet_res_stack is not None:
-                x = x + controlnet_res_stack[block_id]
+        for block in dit.blocks:
+            x = block(x, context, t_mod, freqs)
         if tea_cache is not None:
             tea_cache.store(x)
 

diffsynth/vram_management/layers.py

@@ -70,6 +70,56 @@ class AutoWrappedLinear(torch.nn.Linear):
             bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device)
         return torch.nn.functional.linear(x, weight, bias)
 
+class AutoLoRALinear(torch.nn.Linear):
+    def __init__(self, name='', in_features=1, out_features=2, bias=True, device=None, dtype=None):
+        super().__init__(in_features, out_features, bias, device, dtype)
+        self.name = name
+    
+    def forward(self, x, lora_state_dicts=[], lora_alphas=[1.0,1.0], lora_patcher=None, **kwargs):
+        out = torch.nn.functional.linear(x, self.weight, self.bias)
+        lora_a_name = f'{self.name}.lora_A.default.weight'
+        lora_b_name = f'{self.name}.lora_B.default.weight'
+
+        lora_output = []
+        for i, lora_state_dict in enumerate(lora_state_dicts):
+            if lora_state_dict is None:
+                break
+            if lora_a_name in lora_state_dict and lora_b_name in lora_state_dict:
+                lora_A = lora_state_dict[lora_a_name].to(dtype=self.weight.dtype,device=self.weight.device)
+                lora_B = lora_state_dict[lora_b_name].to(dtype=self.weight.dtype,device=self.weight.device)
+                out_lora = x @ lora_A.T @ lora_B.T
+                lora_output.append(out_lora)
+        if len(lora_output) > 0:
+            lora_output = torch.stack(lora_output)
+            out = lora_patcher(out, lora_output, self.name)
+        return out
+
+def enable_auto_lora(model:torch.nn.Module, module_map: dict, name_prefix=''):
+    targets = list(module_map.keys())
+    for name, module in model.named_children():
+        if name_prefix != '':
+            full_name = name_prefix + '.' + name
+        else:
+            full_name = name
+        if isinstance(module,targets[1]):
+            # print(full_name)
+            # print(module)
+            # ToDo: replace the linear to the AutoLoRALinear 
+            new_module = AutoLoRALinear(
+                name=full_name, 
+                in_features=module.in_features, 
+                out_features=module.out_features, 
+                bias=module.bias is not None, 
+                device=module.weight.device, 
+                dtype=module.weight.dtype)
+            new_module.weight.data.copy_(module.weight.data)
+            new_module.bias.data.copy_(module.bias.data)
+            setattr(model, name, new_module)
+        elif isinstance(module, targets[0]):
+            pass
+        else:
+            enable_auto_lora(module, module_map, full_name)
+       
 
 def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0):
     for name, module in model.named_children():

examples/wanvideo/train_wan_t2v.py

@@ -12,12 +12,9 @@ import numpy as np
 
 
 class TextVideoDataset(torch.utils.data.Dataset):
-    def __init__(self, base_path, metadata_path, max_num_frames=81, frame_interval=1, num_frames=81, height=480, width=832, is_i2v=False, target_fps=None):
+    def __init__(self, base_path, metadata_path, max_num_frames=81, frame_interval=1, num_frames=81, height=480, width=832, is_i2v=False):
         metadata = pd.read_csv(metadata_path)
-        if os.path.exists(os.path.join(base_path, "train")):
-            self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
-        else:
-            self.path = [os.path.join(base_path, file_name) for file_name in metadata["file_name"]]
+        self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
         self.text = metadata["text"].to_list()
         
         self.max_num_frames = max_num_frames
@@ -26,7 +23,6 @@ class TextVideoDataset(torch.utils.data.Dataset):
         self.height = height
         self.width = width
         self.is_i2v = is_i2v
-        self.target_fps = target_fps
             
         self.frame_process = v2.Compose([
             v2.CenterCrop(size=(height, width)),
@@ -75,15 +71,8 @@ class TextVideoDataset(torch.utils.data.Dataset):
 
 
     def load_video(self, file_path):
-        start_frame_id = 0
-        if self.target_fps is None:
-            frame_interval = self.frame_interval
-        else:
-            reader = imageio.get_reader(file_path)
-            fps = reader.get_meta_data()["fps"]
-            reader.close()
-            frame_interval = max(round(fps / self.target_fps), 1)
-        frames = self.load_frames_using_imageio(file_path, self.max_num_frames, start_frame_id, frame_interval, self.num_frames, self.frame_process)
+        start_frame_id = torch.randint(0, self.max_num_frames - (self.num_frames - 1) * self.frame_interval, (1,))[0]
+        frames = self.load_frames_using_imageio(file_path, self.max_num_frames, start_frame_id, self.frame_interval, self.num_frames, self.frame_process)
         return frames
     
     
@@ -106,20 +95,17 @@ class TextVideoDataset(torch.utils.data.Dataset):
     def __getitem__(self, data_id):
         text = self.text[data_id]
         path = self.path[data_id]
-        try:
-            if self.is_image(path):
-                if self.is_i2v:
-                    raise ValueError(f"{path} is not a video. I2V model doesn't support image-to-image training.")
-                video = self.load_image(path)
-            else:
-                video = self.load_video(path)
+        if self.is_image(path):
             if self.is_i2v:
-                video, first_frame = video
-                data = {"text": text, "video": video, "path": path, "first_frame": first_frame}
-            else:
-                data = {"text": text, "video": video, "path": path}
-        except:
-            data = None
+                raise ValueError(f"{path} is not a video. I2V model doesn't support image-to-image training.")
+            video = self.load_image(path)
+        else:
+            video = self.load_video(path)
+        if self.is_i2v:
+            video, first_frame = video
+            data = {"text": text, "video": video, "path": path, "first_frame": first_frame}
+        else:
+            data = {"text": text, "video": video, "path": path}
         return data
     
 
@@ -129,7 +115,7 @@ class TextVideoDataset(torch.utils.data.Dataset):
 
 
 class LightningModelForDataProcess(pl.LightningModule):
-    def __init__(self, text_encoder_path, vae_path, image_encoder_path=None, tiled=False, tile_size=(34, 34), tile_stride=(18, 16), redirected_tensor_path=None):
+    def __init__(self, text_encoder_path, vae_path, image_encoder_path=None, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)):
         super().__init__()
         model_path = [text_encoder_path, vae_path]
         if image_encoder_path is not None:
@@ -139,13 +125,9 @@ class LightningModelForDataProcess(pl.LightningModule):
         self.pipe = WanVideoPipeline.from_model_manager(model_manager)
 
         self.tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}
-        self.redirected_tensor_path = redirected_tensor_path
         
     def test_step(self, batch, batch_idx):
-        data = batch[0]
-        if data is None or data["video"] is None:
-            return
-        text, video, path = data["text"], data["video"].unsqueeze(0), data["path"]
+        text, video, path = batch["text"][0], batch["video"], batch["path"][0]
         
         self.pipe.device = self.device
         if video is not None:
@@ -162,49 +144,28 @@ class LightningModelForDataProcess(pl.LightningModule):
             else:
                 image_emb = {}
             data = {"latents": latents, "prompt_emb": prompt_emb, "image_emb": image_emb}
-            if self.redirected_tensor_path is not None:
-                path = path.replace("/", "_").replace("\\", "_")
-                path = os.path.join(self.redirected_tensor_path, path)
             torch.save(data, path + ".tensors.pth")
 
 
 
 class TensorDataset(torch.utils.data.Dataset):
-    def __init__(self, base_path, metadata_path=None, steps_per_epoch=1000, redirected_tensor_path=None):
-        if os.path.exists(metadata_path):
-            metadata = pd.read_csv(metadata_path)
-            self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
-            print(len(self.path), "videos in metadata.")
-            if redirected_tensor_path is None:
-                self.path = [i + ".tensors.pth" for i in self.path if os.path.exists(i + ".tensors.pth")]
-            else:
-                cached_path = []
-                for path in self.path:
-                    path = path.replace("/", "_").replace("\\", "_")
-                    path = os.path.join(redirected_tensor_path, path)
-                    if os.path.exists(path + ".tensors.pth"):
-                        cached_path.append(path + ".tensors.pth")
-                self.path = cached_path
-        else:
-            print("Cannot find metadata.csv. Trying to search for tensor files.")
-            self.path = [os.path.join(base_path, i) for i in os.listdir(base_path) if i.endswith(".tensors.pth")]
+    def __init__(self, base_path, metadata_path, steps_per_epoch):
+        metadata = pd.read_csv(metadata_path)
+        self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
+        print(len(self.path), "videos in metadata.")
+        self.path = [i + ".tensors.pth" for i in self.path if os.path.exists(i + ".tensors.pth")]
         print(len(self.path), "tensors cached in metadata.")
         assert len(self.path) > 0
         
         self.steps_per_epoch = steps_per_epoch
-        self.redirected_tensor_path = redirected_tensor_path
 
 
     def __getitem__(self, index):
-        while True:
-            try:
-                data_id = torch.randint(0, len(self.path), (1,))[0]
-                data_id = (data_id + index) % len(self.path) # For fixed seed.
-                path = self.path[data_id]
-                data = torch.load(path, weights_only=True, map_location="cpu")
-                return data
-            except:
-                continue
+        data_id = torch.randint(0, len(self.path), (1,))[0]
+        data_id = (data_id + index) % len(self.path) # For fixed seed.
+        path = self.path[data_id]
+        data = torch.load(path, weights_only=True, map_location="cpu")
+        return data
     
 
     def __len__(self):
@@ -362,18 +323,6 @@ def parse_args():
         default="./",
         help="Path to save the model.",
     )
-    parser.add_argument(
-        "--metadata_path",
-        type=str,
-        default=None,
-        help="Path to metadata.csv.",
-    )
-    parser.add_argument(
-        "--redirected_tensor_path",
-        type=str,
-        default=None,
-        help="Path to save cached tensors.",
-    )
     parser.add_argument(
         "--text_encoder_path",
         type=str,
@@ -440,12 +389,6 @@ def parse_args():
         default=81,
         help="Number of frames.",
     )
-    parser.add_argument(
-        "--target_fps",
-        type=int,
-        default=None,
-        help="Expected FPS for sampling frames.",
-    )
     parser.add_argument(
         "--height",
         type=int,
@@ -557,21 +500,19 @@ def parse_args():
 def data_process(args):
     dataset = TextVideoDataset(
         args.dataset_path,
-        os.path.join(args.dataset_path, "metadata.csv") if args.metadata_path is None else args.metadata_path,
+        os.path.join(args.dataset_path, "metadata.csv"),
         max_num_frames=args.num_frames,
         frame_interval=1,
         num_frames=args.num_frames,
         height=args.height,
         width=args.width,
-        is_i2v=args.image_encoder_path is not None,
-        target_fps=args.target_fps,
+        is_i2v=args.image_encoder_path is not None
     )
     dataloader = torch.utils.data.DataLoader(
         dataset,
         shuffle=False,
         batch_size=1,
-        num_workers=args.dataloader_num_workers,
-        collate_fn=lambda x: x,
+        num_workers=args.dataloader_num_workers
     )
     model = LightningModelForDataProcess(
         text_encoder_path=args.text_encoder_path,
@@ -580,7 +521,6 @@ def data_process(args):
         tiled=args.tiled,
         tile_size=(args.tile_size_height, args.tile_size_width),
         tile_stride=(args.tile_stride_height, args.tile_stride_width),
-        redirected_tensor_path=args.redirected_tensor_path,
     )
     trainer = pl.Trainer(
         accelerator="gpu",
@@ -593,9 +533,8 @@ def data_process(args):
 def train(args):
     dataset = TensorDataset(
         args.dataset_path,
-        os.path.join(args.dataset_path, "metadata.csv") if args.metadata_path is None else args.metadata_path,
+        os.path.join(args.dataset_path, "metadata.csv"),
         steps_per_epoch=args.steps_per_epoch,
-        redirected_tensor_path=args.redirected_tensor_path,
     )
     dataloader = torch.utils.data.DataLoader(
         dataset,

examples/wanvideo/train_wan_t2v_controlnet.py

@@ -1,626 +0,0 @@
-import torch, os, imageio, argparse
-from torchvision.transforms import v2
-from einops import rearrange
-import lightning as pl
-import pandas as pd
-from diffsynth import WanVideoPipeline, ModelManager, load_state_dict
-from peft import LoraConfig, inject_adapter_in_model
-import torchvision
-from PIL import Image
-import numpy as np
-from diffsynth.models.wan_video_controlnet import WanControlNetModel
-from diffsynth.pipelines.wan_video import model_fn_wan_video
-
-
-
-class TextVideoDataset(torch.utils.data.Dataset):
-    def __init__(self, base_path, metadata_path, max_num_frames=81, frame_interval=1, num_frames=81, height=480, width=832, is_i2v=False, target_fps=None):
-        metadata = pd.read_csv(metadata_path)
-        self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
-        self.controlnet_path = [os.path.join(base_path, file_name) for file_name in metadata["controlnet_file_name"]]
-        self.text = metadata["text"].to_list()
-        
-        self.max_num_frames = max_num_frames
-        self.frame_interval = frame_interval
-        self.num_frames = num_frames
-        self.height = height
-        self.width = width
-        self.is_i2v = is_i2v
-        self.target_fps = target_fps
-            
-        self.frame_process = v2.Compose([
-            v2.CenterCrop(size=(height, width)),
-            v2.Resize(size=(height, width), antialias=True),
-            v2.ToTensor(),
-            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ])
-        
-        
-    def crop_and_resize(self, image):
-        width, height = image.size
-        scale = max(self.width / width, self.height / height)
-        image = torchvision.transforms.functional.resize(
-            image,
-            (round(height*scale), round(width*scale)),
-            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
-        )
-        return image
-
-
-    def load_frames_using_imageio(self, file_path, max_num_frames, start_frame_id, interval, num_frames, frame_process):
-        reader = imageio.get_reader(file_path)
-        if reader.count_frames() < max_num_frames or reader.count_frames() - 1 < start_frame_id + (num_frames - 1) * interval:
-            reader.close()
-            return None
-        
-        frames = []
-        first_frame = None
-        for frame_id in range(num_frames):
-            frame = reader.get_data(start_frame_id + frame_id * interval)
-            frame = Image.fromarray(frame)
-            frame = self.crop_and_resize(frame)
-            if first_frame is None:
-                first_frame = np.array(frame)
-            frame = frame_process(frame)
-            frames.append(frame)
-        reader.close()
-
-        frames = torch.stack(frames, dim=0)
-        frames = rearrange(frames, "T C H W -> C T H W")
-
-        if self.is_i2v:
-            return frames, first_frame
-        else:
-            return frames
-
-
-    def load_video(self, file_path):
-        start_frame_id = 0
-        if self.target_fps is None:
-            frame_interval = self.frame_interval
-        else:
-            reader = imageio.get_reader(file_path)
-            fps = reader.get_meta_data()["fps"]
-            reader.close()
-            frame_interval = max(round(fps / self.target_fps), 1)
-        frames = self.load_frames_using_imageio(file_path, self.max_num_frames, start_frame_id, frame_interval, self.num_frames, self.frame_process)
-        return frames
-    
-    
-    def is_image(self, file_path):
-        file_ext_name = file_path.split(".")[-1]
-        if file_ext_name.lower() in ["jpg", "jpeg", "png", "webp"]:
-            return True
-        return False
-    
-    
-    def load_image(self, file_path):
-        frame = Image.open(file_path).convert("RGB")
-        frame = self.crop_and_resize(frame)
-        frame = self.frame_process(frame)
-        frame = rearrange(frame, "C H W -> C 1 H W")
-        return frame
-
-
-    def __getitem__(self, data_id):
-        text = self.text[data_id]
-        path = self.path[data_id]
-        controlnet_path = self.controlnet_path[data_id]
-        try:
-            if self.is_image(path):
-                if self.is_i2v:
-                    raise ValueError(f"{path} is not a video. I2V model doesn't support image-to-image training.")
-                video = self.load_image(path)
-            else:
-                video = self.load_video(path)
-                controlnet_frames = self.load_video(controlnet_path)
-            if self.is_i2v:
-                video, first_frame = video
-                data = {"text": text, "video": video, "path": path, "first_frame": first_frame}
-            else:
-                data = {"text": text, "video": video, "path": path, "controlnet_frames": controlnet_frames}
-        except:
-            data = None
-        return data
-    
-
-    def __len__(self):
-        return len(self.path)
-
-
-
-class LightningModelForDataProcess(pl.LightningModule):
-    def __init__(self, text_encoder_path, vae_path, image_encoder_path=None, tiled=False, tile_size=(34, 34), tile_stride=(18, 16), redirected_tensor_path=None):
-        super().__init__()
-        model_path = [text_encoder_path, vae_path]
-        if image_encoder_path is not None:
-            model_path.append(image_encoder_path)
-        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
-        model_manager.load_models(model_path)
-        self.pipe = WanVideoPipeline.from_model_manager(model_manager)
-
-        self.tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}
-        self.redirected_tensor_path = redirected_tensor_path
-        
-    def test_step(self, batch, batch_idx):
-        data = batch[0]
-        if data is None or data["video"] is None:
-            return
-        text, video, path = data["text"], data["video"].unsqueeze(0), data["path"]
-        controlnet_frames = data["controlnet_frames"].unsqueeze(0)
-        
-        self.pipe.device = self.device
-        if video is not None:
-            # prompt
-            prompt_emb = self.pipe.encode_prompt(text)
-            # video
-            video = video.to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
-            latents = self.pipe.encode_video(video, **self.tiler_kwargs)[0]
-            # ControlNet video
-            controlnet_frames = controlnet_frames.to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
-            controlnet_kwargs = self.pipe.prepare_controlnet(controlnet_frames, **self.tiler_kwargs)
-            controlnet_kwargs["controlnet_conditioning"] = controlnet_kwargs["controlnet_conditioning"][0]
-            # image
-            if "first_frame" in batch:
-                first_frame = Image.fromarray(batch["first_frame"][0].cpu().numpy())
-                _, _, num_frames, height, width = video.shape
-                image_emb = self.pipe.encode_image(first_frame, num_frames, height, width)
-            else:
-                image_emb = {}
-            data = {"latents": latents, "prompt_emb": prompt_emb, "image_emb": image_emb, "controlnet_kwargs": controlnet_kwargs}
-            if self.redirected_tensor_path is not None:
-                path = path.replace("/", "_").replace("\\", "_")
-                path = os.path.join(self.redirected_tensor_path, path)
-            torch.save(data, path + ".tensors.pth")
-
-
-
-class TensorDataset(torch.utils.data.Dataset):
-    def __init__(self, base_path, metadata_path=None, steps_per_epoch=1000, redirected_tensor_path=None):
-        if os.path.exists(metadata_path):
-            metadata = pd.read_csv(metadata_path)
-            self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
-            print(len(self.path), "videos in metadata.")
-            if redirected_tensor_path is None:
-                self.path = [i + ".tensors.pth" for i in self.path if os.path.exists(i + ".tensors.pth")]
-            else:
-                cached_path = []
-                for path in self.path:
-                    path = path.replace("/", "_").replace("\\", "_")
-                    path = os.path.join(redirected_tensor_path, path)
-                    if os.path.exists(path + ".tensors.pth"):
-                        cached_path.append(path + ".tensors.pth")
-                self.path = cached_path
-        else:
-            print("Cannot find metadata.csv. Trying to search for tensor files.")
-            self.path = [os.path.join(base_path, i) for i in os.listdir(base_path) if i.endswith(".tensors.pth")]
-        print(len(self.path), "tensors cached in metadata.")
-        assert len(self.path) > 0
-        
-        self.steps_per_epoch = steps_per_epoch
-        self.redirected_tensor_path = redirected_tensor_path
-
-
-    def __getitem__(self, index):
-        while True:
-            try:
-                data_id = torch.randint(0, len(self.path), (1,))[0]
-                data_id = (data_id + index) % len(self.path) # For fixed seed.
-                path = self.path[data_id]
-                data = torch.load(path, weights_only=True, map_location="cpu")
-                return data
-            except:
-                continue
-    
-
-    def __len__(self):
-        return self.steps_per_epoch
-
-
-
-class LightningModelForTrain(pl.LightningModule):
-    def __init__(
-        self,
-        dit_path,
-        learning_rate=1e-5,
-        lora_rank=4, lora_alpha=4, train_architecture="lora", lora_target_modules="q,k,v,o,ffn.0,ffn.2", init_lora_weights="kaiming",
-        use_gradient_checkpointing=True, use_gradient_checkpointing_offload=False,
-        pretrained_lora_path=None
-    ):
-        super().__init__()
-        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
-        if os.path.isfile(dit_path):
-            model_manager.load_models([dit_path])
-        else:
-            dit_path = dit_path.split(",")
-            model_manager.load_models([dit_path])
-        
-        self.pipe = WanVideoPipeline.from_model_manager(model_manager)
-        self.pipe.scheduler.set_timesteps(1000, training=True)
-        self.freeze_parameters()
-        
-        state_dict = load_state_dict(dit_path, torch_dtype=torch.bfloat16)
-        state_dict, config = WanControlNetModel.state_dict_converter().from_base_model(state_dict)
-        self.pipe.controlnet = WanControlNetModel(**config).to(torch.bfloat16)
-        self.pipe.controlnet.load_state_dict(state_dict)
-        self.pipe.controlnet.train()
-        self.pipe.controlnet.requires_grad_(True)
-        
-        self.learning_rate = learning_rate
-        self.use_gradient_checkpointing = use_gradient_checkpointing
-        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
-        
-        
-    def freeze_parameters(self):
-        # Freeze parameters
-        self.pipe.requires_grad_(False)
-        self.pipe.eval()
-        self.pipe.denoising_model().train()
-
-
-    def training_step(self, batch, batch_idx):
-        # Data
-        latents = batch["latents"].to(self.device)
-        controlnet_kwargs = batch["controlnet_kwargs"]
-        controlnet_kwargs["controlnet_conditioning"] = controlnet_kwargs["controlnet_conditioning"].to(self.device)
-        prompt_emb = batch["prompt_emb"]
-        prompt_emb["context"] = prompt_emb["context"][0].to(self.device)
-        image_emb = batch["image_emb"]
-        if "clip_feature" in image_emb:
-            image_emb["clip_feature"] = image_emb["clip_feature"][0].to(self.device)
-        if "y" in image_emb:
-            image_emb["y"] = image_emb["y"][0].to(self.device)
-
-        # Loss
-        self.pipe.device = self.device
-        noise = torch.randn_like(latents)
-        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))
-        timestep = self.pipe.scheduler.timesteps[timestep_id].to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
-        extra_input = self.pipe.prepare_extra_input(latents)
-        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)
-        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)
-
-        # Compute loss
-        noise_pred = model_fn_wan_video(
-            dit=self.pipe.dit, controlnet=self.pipe.controlnet,
-            x=noisy_latents, timestep=timestep, **prompt_emb, **extra_input, **image_emb, **controlnet_kwargs,
-            use_gradient_checkpointing=self.use_gradient_checkpointing,
-            use_gradient_checkpointing_offload=self.use_gradient_checkpointing_offload
-        )
-        loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
-        loss = loss * self.pipe.scheduler.training_weight(timestep)
-
-        # Record log
-        self.log("train_loss", loss, prog_bar=True)
-        return loss
-
-
-    def configure_optimizers(self):
-        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.controlnet.parameters())
-        optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate)
-        return optimizer
-    
-
-    def on_save_checkpoint(self, checkpoint):
-        checkpoint.clear()
-        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.pipe.controlnet.named_parameters()))
-        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
-        state_dict = self.pipe.controlnet.state_dict()
-        lora_state_dict = {}
-        for name, param in state_dict.items():
-            if name in trainable_param_names:
-                lora_state_dict[name] = param
-        checkpoint.update(lora_state_dict)
-
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Simple example of a training script.")
-    parser.add_argument(
-        "--task",
-        type=str,
-        default="data_process",
-        required=True,
-        choices=["data_process", "train"],
-        help="Task. `data_process` or `train`.",
-    )
-    parser.add_argument(
-        "--dataset_path",
-        type=str,
-        default=None,
-        required=True,
-        help="The path of the Dataset.",
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        default="./",
-        help="Path to save the model.",
-    )
-    parser.add_argument(
-        "--metadata_path",
-        type=str,
-        default=None,
-        help="Path to metadata.csv.",
-    )
-    parser.add_argument(
-        "--redirected_tensor_path",
-        type=str,
-        default=None,
-        help="Path to save cached tensors.",
-    )
-    parser.add_argument(
-        "--text_encoder_path",
-        type=str,
-        default=None,
-        help="Path of text encoder.",
-    )
-    parser.add_argument(
-        "--image_encoder_path",
-        type=str,
-        default=None,
-        help="Path of image encoder.",
-    )
-    parser.add_argument(
-        "--vae_path",
-        type=str,
-        default=None,
-        help="Path of VAE.",
-    )
-    parser.add_argument(
-        "--dit_path",
-        type=str,
-        default=None,
-        help="Path of DiT.",
-    )
-    parser.add_argument(
-        "--tiled",
-        default=False,
-        action="store_true",
-        help="Whether enable tile encode in VAE. This option can reduce VRAM required.",
-    )
-    parser.add_argument(
-        "--tile_size_height",
-        type=int,
-        default=34,
-        help="Tile size (height) in VAE.",
-    )
-    parser.add_argument(
-        "--tile_size_width",
-        type=int,
-        default=34,
-        help="Tile size (width) in VAE.",
-    )
-    parser.add_argument(
-        "--tile_stride_height",
-        type=int,
-        default=18,
-        help="Tile stride (height) in VAE.",
-    )
-    parser.add_argument(
-        "--tile_stride_width",
-        type=int,
-        default=16,
-        help="Tile stride (width) in VAE.",
-    )
-    parser.add_argument(
-        "--steps_per_epoch",
-        type=int,
-        default=500,
-        help="Number of steps per epoch.",
-    )
-    parser.add_argument(
-        "--num_frames",
-        type=int,
-        default=81,
-        help="Number of frames.",
-    )
-    parser.add_argument(
-        "--target_fps",
-        type=int,
-        default=None,
-        help="Expected FPS for sampling frames.",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=480,
-        help="Image height.",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=832,
-        help="Image width.",
-    )
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=1,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=1e-5,
-        help="Learning rate.",
-    )
-    parser.add_argument(
-        "--accumulate_grad_batches",
-        type=int,
-        default=1,
-        help="The number of batches in gradient accumulation.",
-    )
-    parser.add_argument(
-        "--max_epochs",
-        type=int,
-        default=1,
-        help="Number of epochs.",
-    )
-    parser.add_argument(
-        "--lora_target_modules",
-        type=str,
-        default="q,k,v,o,ffn.0,ffn.2",
-        help="Layers with LoRA modules.",
-    )
-    parser.add_argument(
-        "--init_lora_weights",
-        type=str,
-        default="kaiming",
-        choices=["gaussian", "kaiming"],
-        help="The initializing method of LoRA weight.",
-    )
-    parser.add_argument(
-        "--training_strategy",
-        type=str,
-        default="auto",
-        choices=["auto", "deepspeed_stage_1", "deepspeed_stage_2", "deepspeed_stage_3"],
-        help="Training strategy",
-    )
-    parser.add_argument(
-        "--lora_rank",
-        type=int,
-        default=4,
-        help="The dimension of the LoRA update matrices.",
-    )
-    parser.add_argument(
-        "--lora_alpha",
-        type=float,
-        default=4.0,
-        help="The weight of the LoRA update matrices.",
-    )
-    parser.add_argument(
-        "--use_gradient_checkpointing",
-        default=False,
-        action="store_true",
-        help="Whether to use gradient checkpointing.",
-    )
-    parser.add_argument(
-        "--use_gradient_checkpointing_offload",
-        default=False,
-        action="store_true",
-        help="Whether to use gradient checkpointing offload.",
-    )
-    parser.add_argument(
-        "--train_architecture",
-        type=str,
-        default="lora",
-        choices=["lora", "full"],
-        help="Model structure to train. LoRA training or full training.",
-    )
-    parser.add_argument(
-        "--pretrained_lora_path",
-        type=str,
-        default=None,
-        help="Pretrained LoRA path. Required if the training is resumed.",
-    )
-    parser.add_argument(
-        "--use_swanlab",
-        default=False,
-        action="store_true",
-        help="Whether to use SwanLab logger.",
-    )
-    parser.add_argument(
-        "--swanlab_mode",
-        default=None,
-        help="SwanLab mode (cloud or local).",
-    )
-    args = parser.parse_args()
-    return args
-
-
-def data_process(args):
-    dataset = TextVideoDataset(
-        args.dataset_path,
-        os.path.join(args.dataset_path, "metadata.csv") if args.metadata_path is None else args.metadata_path,
-        max_num_frames=args.num_frames,
-        frame_interval=1,
-        num_frames=args.num_frames,
-        height=args.height,
-        width=args.width,
-        is_i2v=args.image_encoder_path is not None,
-        target_fps=args.target_fps,
-    )
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        shuffle=False,
-        batch_size=1,
-        num_workers=args.dataloader_num_workers,
-        collate_fn=lambda x: x,
-    )
-    model = LightningModelForDataProcess(
-        text_encoder_path=args.text_encoder_path,
-        image_encoder_path=args.image_encoder_path,
-        vae_path=args.vae_path,
-        tiled=args.tiled,
-        tile_size=(args.tile_size_height, args.tile_size_width),
-        tile_stride=(args.tile_stride_height, args.tile_stride_width),
-        redirected_tensor_path=args.redirected_tensor_path,
-    )
-    trainer = pl.Trainer(
-        accelerator="gpu",
-        devices="auto",
-        default_root_dir=args.output_path,
-    )
-    trainer.test(model, dataloader)
-    
-    
-def train(args):
-    dataset = TensorDataset(
-        args.dataset_path,
-        os.path.join(args.dataset_path, "metadata.csv") if args.metadata_path is None else args.metadata_path,
-        steps_per_epoch=args.steps_per_epoch,
-        redirected_tensor_path=args.redirected_tensor_path,
-    )
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        shuffle=True,
-        batch_size=1,
-        num_workers=args.dataloader_num_workers
-    )
-    model = LightningModelForTrain(
-        dit_path=args.dit_path,
-        learning_rate=args.learning_rate,
-        train_architecture=args.train_architecture,
-        lora_rank=args.lora_rank,
-        lora_alpha=args.lora_alpha,
-        lora_target_modules=args.lora_target_modules,
-        init_lora_weights=args.init_lora_weights,
-        use_gradient_checkpointing=args.use_gradient_checkpointing,
-        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
-        pretrained_lora_path=args.pretrained_lora_path,
-    )
-    if args.use_swanlab:
-        from swanlab.integration.pytorch_lightning import SwanLabLogger
-        swanlab_config = {"UPPERFRAMEWORK": "DiffSynth-Studio"}
-        swanlab_config.update(vars(args))
-        swanlab_logger = SwanLabLogger(
-            project="wan", 
-            name="wan",
-            config=swanlab_config,
-            mode=args.swanlab_mode,
-            logdir=os.path.join(args.output_path, "swanlog"),
-        )
-        logger = [swanlab_logger]
-    else:
-        logger = None
-    trainer = pl.Trainer(
-        max_epochs=args.max_epochs,
-        accelerator="gpu",
-        devices="auto",
-        precision="bf16",
-        strategy=args.training_strategy,
-        default_root_dir=args.output_path,
-        accumulate_grad_batches=args.accumulate_grad_batches,
-        callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)],
-        logger=logger,
-    )
-    trainer.fit(model, dataloader)
-
-
-if __name__ == '__main__':
-    args = parse_args()
-    if args.task == "data_process":
-        data_process(args)
-    elif args.task == "train":
-        train(args)

examples/wanvideo/train_wan_t2v_motion.py

@@ -1,691 +0,0 @@
-import torch, os, imageio, argparse
-from torchvision.transforms import v2
-from einops import rearrange
-import lightning as pl
-import pandas as pd
-from diffsynth import WanVideoPipeline, ModelManager, load_state_dict
-from diffsynth.models.wan_video_motion_controller import WanMotionControllerModel
-from diffsynth.pipelines.wan_video import model_fn_wan_video
-from peft import LoraConfig, inject_adapter_in_model
-import torchvision
-from PIL import Image
-import numpy as np
-from tqdm import tqdm
-
-
-
-class TextVideoDataset(torch.utils.data.Dataset):
-    def __init__(self, base_path, metadata_path, max_num_frames=81, frame_interval=1, num_frames=81, height=480, width=832, is_i2v=False, target_fps=None):
-        metadata = pd.read_csv(metadata_path)
-        self.path = [os.path.join(base_path, file_name) for file_name in metadata["file_name"]]
-        self.text = metadata["text"].to_list()
-        
-        self.max_num_frames = max_num_frames
-        self.frame_interval = frame_interval
-        self.num_frames = num_frames
-        self.height = height
-        self.width = width
-        self.is_i2v = is_i2v
-        self.target_fps = target_fps
-            
-        self.frame_process = v2.Compose([
-            v2.CenterCrop(size=(height, width)),
-            v2.Resize(size=(height, width), antialias=True),
-            v2.ToTensor(),
-            v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ])
-        
-        
-    def crop_and_resize(self, image):
-        width, height = image.size
-        scale = max(self.width / width, self.height / height)
-        image = torchvision.transforms.functional.resize(
-            image,
-            (round(height*scale), round(width*scale)),
-            interpolation=torchvision.transforms.InterpolationMode.BILINEAR
-        )
-        return image
-
-
-    def load_frames_using_imageio(self, file_path, max_num_frames, start_frame_id, interval, num_frames, frame_process):
-        reader = imageio.get_reader(file_path)
-        if reader.count_frames() < max_num_frames or reader.count_frames() - 1 < start_frame_id + (num_frames - 1) * interval:
-            reader.close()
-            return None
-        
-        frames = []
-        first_frame = None
-        for frame_id in range(num_frames):
-            frame = reader.get_data(start_frame_id + frame_id * interval)
-            frame = Image.fromarray(frame)
-            frame = self.crop_and_resize(frame)
-            if first_frame is None:
-                first_frame = np.array(frame)
-            frame = frame_process(frame)
-            frames.append(frame)
-        reader.close()
-
-        frames = torch.stack(frames, dim=0)
-        frames = rearrange(frames, "T C H W -> C T H W")
-
-        if self.is_i2v:
-            return frames, first_frame
-        else:
-            return frames
-
-
-    def load_video(self, file_path):
-        start_frame_id = 0
-        if self.target_fps is None:
-            frame_interval = self.frame_interval
-        else:
-            reader = imageio.get_reader(file_path)
-            fps = reader.get_meta_data()["fps"]
-            reader.close()
-            frame_interval = max(round(fps / self.target_fps), 1)
-        frames = self.load_frames_using_imageio(file_path, self.max_num_frames, start_frame_id, frame_interval, self.num_frames, self.frame_process)
-        return frames
-    
-    
-    def is_image(self, file_path):
-        file_ext_name = file_path.split(".")[-1]
-        if file_ext_name.lower() in ["jpg", "jpeg", "png", "webp"]:
-            return True
-        return False
-    
-    
-    def load_image(self, file_path):
-        frame = Image.open(file_path).convert("RGB")
-        frame = self.crop_and_resize(frame)
-        first_frame = frame
-        frame = self.frame_process(frame)
-        frame = rearrange(frame, "C H W -> C 1 H W")
-        return frame
-
-
-    def __getitem__(self, data_id):
-        text = self.text[data_id]
-        path = self.path[data_id]
-        try:
-            if self.is_image(path):
-                if self.is_i2v:
-                    raise ValueError(f"{path} is not a video. I2V model doesn't support image-to-image training.")
-                video = self.load_image(path)
-            else:
-                video = self.load_video(path)
-            if self.is_i2v:
-                video, first_frame = video
-                data = {"text": text, "video": video, "path": path, "first_frame": first_frame}
-            else:
-                data = {"text": text, "video": video, "path": path}
-        except:
-            data = None
-        return data
-    
-
-    def __len__(self):
-        return len(self.path)
-
-
-
-class LightningModelForDataProcess(pl.LightningModule):
-    def __init__(self, text_encoder_path, vae_path, image_encoder_path=None, tiled=False, tile_size=(34, 34), tile_stride=(18, 16), redirected_tensor_path=None):
-        super().__init__()
-        model_path = [text_encoder_path, vae_path]
-        if image_encoder_path is not None:
-            model_path.append(image_encoder_path)
-        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
-        model_manager.load_models(model_path)
-        self.pipe = WanVideoPipeline.from_model_manager(model_manager)
-
-        self.tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}
-        self.redirected_tensor_path = redirected_tensor_path
-        
-    def test_step(self, batch, batch_idx):
-        data = batch[0]
-        if data is None or data["video"] is None:
-            return
-        text, video, path = data["text"], data["video"].unsqueeze(0), data["path"]
-        
-        self.pipe.device = self.device
-        if video is not None:
-            # prompt
-            prompt_emb = self.pipe.encode_prompt(text)
-            # video
-            video = video.to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
-            latents = self.pipe.encode_video(video, **self.tiler_kwargs)[0]
-            # image
-            if "first_frame" in batch:
-                first_frame = Image.fromarray(batch["first_frame"][0].cpu().numpy())
-                _, _, num_frames, height, width = video.shape
-                image_emb = self.pipe.encode_image(first_frame, num_frames, height, width)
-            else:
-                image_emb = {}
-            data = {"latents": latents, "prompt_emb": prompt_emb, "image_emb": image_emb}
-            if self.redirected_tensor_path is not None:
-                path = path.replace("/", "_").replace("\\", "_")
-                path = os.path.join(self.redirected_tensor_path, path)
-            torch.save(data, path + ".tensors.pth")
-
-
-
-class TensorDataset(torch.utils.data.Dataset):
-    def __init__(self, base_path, metadata_path=None, steps_per_epoch=1000, redirected_tensor_path=None):
-        if os.path.exists(metadata_path):
-            metadata = pd.read_csv(metadata_path)
-            self.path = [os.path.join(base_path, "train", file_name) for file_name in metadata["file_name"]]
-            print(len(self.path), "videos in metadata.")
-            if redirected_tensor_path is None:
-                self.path = [i + ".tensors.pth" for i in self.path if os.path.exists(i + ".tensors.pth")]
-            else:
-                cached_path = []
-                for path in self.path:
-                    path = path.replace("/", "_").replace("\\", "_")
-                    path = os.path.join(redirected_tensor_path, path)
-                    if os.path.exists(path + ".tensors.pth"):
-                        cached_path.append(path + ".tensors.pth")
-                self.path = cached_path
-        else:
-            print("Cannot find metadata.csv. Trying to search for tensor files.")
-            self.path = [os.path.join(base_path, i) for i in os.listdir(base_path) if i.endswith(".tensors.pth")]
-        print(len(self.path), "tensors cached in metadata.")
-        assert len(self.path) > 0
-        
-        self.steps_per_epoch = steps_per_epoch
-        self.redirected_tensor_path = redirected_tensor_path
-
-
-    def __getitem__(self, index):
-        while True:
-            try:
-                data_id = torch.randint(0, len(self.path), (1,))[0]
-                data_id = (data_id + index) % len(self.path) # For fixed seed.
-                path = self.path[data_id]
-                data = torch.load(path, weights_only=True, map_location="cpu")
-                return data
-            except:
-                continue
-    
-
-    def __len__(self):
-        return self.steps_per_epoch
-
-
-
-class LightningModelForTrain(pl.LightningModule):
-    def __init__(
-        self,
-        dit_path,
-        learning_rate=1e-5,
-        lora_rank=4, lora_alpha=4, train_architecture="lora", lora_target_modules="q,k,v,o,ffn.0,ffn.2", init_lora_weights="kaiming",
-        use_gradient_checkpointing=True, use_gradient_checkpointing_offload=False,
-        pretrained_lora_path=None
-    ):
-        super().__init__()
-        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
-        if os.path.isfile(dit_path):
-            model_manager.load_models([dit_path])
-        else:
-            dit_path = dit_path.split(",")
-            model_manager.load_models([dit_path])
-        
-        self.pipe = WanVideoPipeline.from_model_manager(model_manager)
-        self.pipe.scheduler.set_timesteps(1000, training=True)
-        self.freeze_parameters()
-        
-        self.pipe.motion_controller = WanMotionControllerModel().to(torch.bfloat16)
-        self.pipe.motion_controller.init()
-        self.pipe.motion_controller.requires_grad_(True)
-        self.pipe.motion_controller.train()
-        self.motion_bucket_manager = MotionBucketManager()
-        
-        self.learning_rate = learning_rate
-        self.use_gradient_checkpointing = use_gradient_checkpointing
-        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
-        
-        
-    def freeze_parameters(self):
-        # Freeze parameters
-        self.pipe.requires_grad_(False)
-        self.pipe.eval()
-        self.pipe.dit.train()
-        
-        
-    def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="q,k,v,o,ffn.0,ffn.2", init_lora_weights="kaiming", pretrained_lora_path=None, state_dict_converter=None):
-        # Add LoRA to UNet
-        self.lora_alpha = lora_alpha
-        if init_lora_weights == "kaiming":
-            init_lora_weights = True
-            
-        lora_config = LoraConfig(
-            r=lora_rank,
-            lora_alpha=lora_alpha,
-            init_lora_weights=init_lora_weights,
-            target_modules=lora_target_modules.split(","),
-        )
-        model = inject_adapter_in_model(lora_config, model)
-        for param in model.parameters():
-            # Upcast LoRA parameters into fp32
-            if param.requires_grad:
-                param.data = param.to(torch.float32)
-                
-        # Lora pretrained lora weights
-        if pretrained_lora_path is not None:
-            state_dict = load_state_dict(pretrained_lora_path)
-            if state_dict_converter is not None:
-                state_dict = state_dict_converter(state_dict)
-            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
-            all_keys = [i for i, _ in model.named_parameters()]
-            num_updated_keys = len(all_keys) - len(missing_keys)
-            num_unexpected_keys = len(unexpected_keys)
-            print(f"{num_updated_keys} parameters are loaded from {pretrained_lora_path}. {num_unexpected_keys} parameters are unexpected.")
-    
-
-    def training_step(self, batch, batch_idx):
-        # Data
-        latents = batch["latents"].to(self.device)
-        prompt_emb = batch["prompt_emb"]
-        prompt_emb["context"] = prompt_emb["context"][0].to(self.device)
-        image_emb = batch["image_emb"]
-        if "clip_feature" in image_emb:
-            image_emb["clip_feature"] = image_emb["clip_feature"][0].to(self.device)
-        if "y" in image_emb:
-            image_emb["y"] = image_emb["y"][0].to(self.device)
-
-        # Loss
-        self.pipe.device = self.device
-        noise = torch.randn_like(latents)
-        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))
-        timestep = self.pipe.scheduler.timesteps[timestep_id].to(dtype=self.pipe.torch_dtype, device=self.pipe.device)
-        extra_input = self.pipe.prepare_extra_input(latents)
-        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)
-        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)
-        motion_bucket_id = self.motion_bucket_manager(latents)
-        motion_bucket_kwargs = self.pipe.prepare_motion_bucket_id(motion_bucket_id)
-
-        # Compute loss
-        noise_pred = model_fn_wan_video(
-            dit=self.pipe.dit, motion_controller=self.pipe.motion_controller,
-            x=noisy_latents, timestep=timestep, **prompt_emb, **extra_input, **image_emb, **motion_bucket_kwargs,
-            use_gradient_checkpointing=self.use_gradient_checkpointing,
-            use_gradient_checkpointing_offload=self.use_gradient_checkpointing_offload
-        )
-        loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
-        loss = loss * self.pipe.scheduler.training_weight(timestep)
-
-        # Record log
-        self.log("train_loss", loss, prog_bar=True)
-        return loss
-
-
-    def configure_optimizers(self):
-        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.motion_controller.parameters())
-        optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate)
-        return optimizer
-    
-
-    def on_save_checkpoint(self, checkpoint):
-        checkpoint.clear()
-        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.pipe.motion_controller.named_parameters()))
-        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
-        state_dict = self.pipe.motion_controller.state_dict()
-        lora_state_dict = {}
-        for name, param in state_dict.items():
-            if name in trainable_param_names:
-                lora_state_dict[name] = param
-        checkpoint.update(lora_state_dict)
-        
-        
-
-class MotionBucketManager:
-    def __init__(self):
-        self.thresholds = [
-            0.093750000, 0.094726562, 0.100585938, 0.100585938, 0.108886719, 0.109375000, 0.118652344, 0.127929688, 0.127929688, 0.130859375,
-            0.133789062, 0.137695312, 0.138671875, 0.138671875, 0.139648438, 0.143554688, 0.143554688, 0.147460938, 0.149414062, 0.149414062,
-            0.152343750, 0.153320312, 0.154296875, 0.154296875, 0.157226562, 0.163085938, 0.163085938, 0.164062500, 0.165039062, 0.166992188,
-            0.173828125, 0.179687500, 0.180664062, 0.184570312, 0.187500000, 0.188476562, 0.188476562, 0.189453125, 0.189453125, 0.202148438,
-            0.206054688, 0.210937500, 0.210937500, 0.211914062, 0.214843750, 0.214843750, 0.216796875, 0.216796875, 0.216796875, 0.218750000,
-            0.218750000, 0.221679688, 0.222656250, 0.227539062, 0.229492188, 0.230468750, 0.236328125, 0.243164062, 0.243164062, 0.245117188,
-            0.253906250, 0.253906250, 0.255859375, 0.259765625, 0.275390625, 0.275390625, 0.277343750, 0.279296875, 0.279296875, 0.279296875,
-            0.292968750, 0.292968750, 0.302734375, 0.306640625, 0.312500000, 0.312500000, 0.326171875, 0.330078125, 0.332031250, 0.332031250,
-            0.337890625, 0.343750000, 0.343750000, 0.351562500, 0.355468750, 0.357421875, 0.361328125, 0.367187500, 0.382812500, 0.388671875,
-            0.392578125, 0.392578125, 0.392578125, 0.404296875, 0.404296875, 0.425781250, 0.433593750, 0.507812500, 0.519531250, 0.539062500,
-        ]
-
-    def get_motion_score(self, frames):
-        score = frames[:, :, 1:, :, :].std(dim=2).mean().tolist()
-        return score
-    
-    def get_bucket_id(self, motion_score):
-        for bucket_id in range(len(self.thresholds) - 1):
-            if self.thresholds[bucket_id + 1] > motion_score:
-                return bucket_id
-        return len(self.thresholds)
-
-    def __call__(self, frames):
-        score = self.get_motion_score(frames)
-        bucket_id = self.get_bucket_id(score)
-        return bucket_id
-
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Simple example of a training script.")
-    parser.add_argument(
-        "--task",
-        type=str,
-        default="data_process",
-        required=True,
-        choices=["data_process", "train"],
-        help="Task. `data_process` or `train`.",
-    )
-    parser.add_argument(
-        "--dataset_path",
-        type=str,
-        default=None,
-        required=True,
-        help="The path of the Dataset.",
-    )
-    parser.add_argument(
-        "--output_path",
-        type=str,
-        default="./",
-        help="Path to save the model.",
-    )
-    parser.add_argument(
-        "--metadata_path",
-        type=str,
-        default=None,
-        help="Path to metadata.csv.",
-    )
-    parser.add_argument(
-        "--redirected_tensor_path",
-        type=str,
-        default=None,
-        help="Path to save cached tensors.",
-    )
-    parser.add_argument(
-        "--text_encoder_path",
-        type=str,
-        default=None,
-        help="Path of text encoder.",
-    )
-    parser.add_argument(
-        "--image_encoder_path",
-        type=str,
-        default=None,
-        help="Path of image encoder.",
-    )
-    parser.add_argument(
-        "--vae_path",
-        type=str,
-        default=None,
-        help="Path of VAE.",
-    )
-    parser.add_argument(
-        "--dit_path",
-        type=str,
-        default=None,
-        help="Path of DiT.",
-    )
-    parser.add_argument(
-        "--tiled",
-        default=False,
-        action="store_true",
-        help="Whether enable tile encode in VAE. This option can reduce VRAM required.",
-    )
-    parser.add_argument(
-        "--tile_size_height",
-        type=int,
-        default=34,
-        help="Tile size (height) in VAE.",
-    )
-    parser.add_argument(
-        "--tile_size_width",
-        type=int,
-        default=34,
-        help="Tile size (width) in VAE.",
-    )
-    parser.add_argument(
-        "--tile_stride_height",
-        type=int,
-        default=18,
-        help="Tile stride (height) in VAE.",
-    )
-    parser.add_argument(
-        "--tile_stride_width",
-        type=int,
-        default=16,
-        help="Tile stride (width) in VAE.",
-    )
-    parser.add_argument(
-        "--steps_per_epoch",
-        type=int,
-        default=500,
-        help="Number of steps per epoch.",
-    )
-    parser.add_argument(
-        "--num_frames",
-        type=int,
-        default=81,
-        help="Number of frames.",
-    )
-    parser.add_argument(
-        "--target_fps",
-        type=int,
-        default=None,
-        help="Expected FPS for sampling frames.",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=480,
-        help="Image height.",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=832,
-        help="Image width.",
-    )
-    parser.add_argument(
-        "--dataloader_num_workers",
-        type=int,
-        default=1,
-        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=1e-5,
-        help="Learning rate.",
-    )
-    parser.add_argument(
-        "--accumulate_grad_batches",
-        type=int,
-        default=1,
-        help="The number of batches in gradient accumulation.",
-    )
-    parser.add_argument(
-        "--max_epochs",
-        type=int,
-        default=1,
-        help="Number of epochs.",
-    )
-    parser.add_argument(
-        "--lora_target_modules",
-        type=str,
-        default="q,k,v,o,ffn.0,ffn.2",
-        help="Layers with LoRA modules.",
-    )
-    parser.add_argument(
-        "--init_lora_weights",
-        type=str,
-        default="kaiming",
-        choices=["gaussian", "kaiming"],
-        help="The initializing method of LoRA weight.",
-    )
-    parser.add_argument(
-        "--training_strategy",
-        type=str,
-        default="auto",
-        choices=["auto", "deepspeed_stage_1", "deepspeed_stage_2", "deepspeed_stage_3"],
-        help="Training strategy",
-    )
-    parser.add_argument(
-        "--lora_rank",
-        type=int,
-        default=4,
-        help="The dimension of the LoRA update matrices.",
-    )
-    parser.add_argument(
-        "--lora_alpha",
-        type=float,
-        default=4.0,
-        help="The weight of the LoRA update matrices.",
-    )
-    parser.add_argument(
-        "--use_gradient_checkpointing",
-        default=False,
-        action="store_true",
-        help="Whether to use gradient checkpointing.",
-    )
-    parser.add_argument(
-        "--use_gradient_checkpointing_offload",
-        default=False,
-        action="store_true",
-        help="Whether to use gradient checkpointing offload.",
-    )
-    parser.add_argument(
-        "--train_architecture",
-        type=str,
-        default="lora",
-        choices=["lora", "full"],
-        help="Model structure to train. LoRA training or full training.",
-    )
-    parser.add_argument(
-        "--pretrained_lora_path",
-        type=str,
-        default=None,
-        help="Pretrained LoRA path. Required if the training is resumed.",
-    )
-    parser.add_argument(
-        "--use_swanlab",
-        default=False,
-        action="store_true",
-        help="Whether to use SwanLab logger.",
-    )
-    parser.add_argument(
-        "--swanlab_mode",
-        default=None,
-        help="SwanLab mode (cloud or local).",
-    )
-    args = parser.parse_args()
-    return args
-
-
-def data_process(args):
-    dataset = TextVideoDataset(
-        args.dataset_path,
-        os.path.join(args.dataset_path, "metadata.csv") if args.metadata_path is None else args.metadata_path,
-        max_num_frames=args.num_frames,
-        frame_interval=1,
-        num_frames=args.num_frames,
-        height=args.height,
-        width=args.width,
-        is_i2v=args.image_encoder_path is not None,
-        target_fps=args.target_fps,
-    )
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        shuffle=False,
-        batch_size=1,
-        num_workers=args.dataloader_num_workers,
-        collate_fn=lambda x: x,
-    )
-    model = LightningModelForDataProcess(
-        text_encoder_path=args.text_encoder_path,
-        image_encoder_path=args.image_encoder_path,
-        vae_path=args.vae_path,
-        tiled=args.tiled,
-        tile_size=(args.tile_size_height, args.tile_size_width),
-        tile_stride=(args.tile_stride_height, args.tile_stride_width),
-        redirected_tensor_path=args.redirected_tensor_path,
-    )
-    trainer = pl.Trainer(
-        accelerator="gpu",
-        devices="auto",
-        default_root_dir=args.output_path,
-    )
-    trainer.test(model, dataloader)
-    
-    
-def get_motion_thresholds(dataloader):
-    scores = []
-    for data in tqdm(dataloader):
-        scores.append(data["latents"][:, :, 1:, :, :].std(dim=2).mean().tolist())
-    scores = sorted(scores)
-    for i in range(100):
-        s = scores[int(i/100 * len(scores))]
-        print("%.9f" % s, end=", ")
-    
-    
-def train(args):
-    dataset = TensorDataset(
-        args.dataset_path,
-        os.path.join(args.dataset_path, "metadata.csv") if args.metadata_path is None else args.metadata_path,
-        steps_per_epoch=args.steps_per_epoch,
-        redirected_tensor_path=args.redirected_tensor_path,
-    )
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        shuffle=True,
-        batch_size=1,
-        num_workers=args.dataloader_num_workers
-    )
-    model = LightningModelForTrain(
-        dit_path=args.dit_path,
-        learning_rate=args.learning_rate,
-        train_architecture=args.train_architecture,
-        lora_rank=args.lora_rank,
-        lora_alpha=args.lora_alpha,
-        lora_target_modules=args.lora_target_modules,
-        init_lora_weights=args.init_lora_weights,
-        use_gradient_checkpointing=args.use_gradient_checkpointing,
-        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
-        pretrained_lora_path=args.pretrained_lora_path,
-    )
-    if args.use_swanlab:
-        from swanlab.integration.pytorch_lightning import SwanLabLogger
-        swanlab_config = {"UPPERFRAMEWORK": "DiffSynth-Studio"}
-        swanlab_config.update(vars(args))
-        swanlab_logger = SwanLabLogger(
-            project="wan", 
-            name="wan",
-            config=swanlab_config,
-            mode=args.swanlab_mode,
-            logdir=os.path.join(args.output_path, "swanlog"),
-        )
-        logger = [swanlab_logger]
-    else:
-        logger = None
-    trainer = pl.Trainer(
-        max_epochs=args.max_epochs,
-        accelerator="gpu",
-        devices="auto",
-        precision="bf16",
-        strategy=args.training_strategy,
-        default_root_dir=args.output_path,
-        accumulate_grad_batches=args.accumulate_grad_batches,
-        callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)],
-        logger=logger,
-    )
-    trainer.fit(model, dataloader)
-
-
-if __name__ == '__main__':
-    args = parse_args()
-    if args.task == "data_process":
-        data_process(args)
-    elif args.task == "train":
-        train(args)

lora/dataset.py

@@ -0,0 +1,54 @@
+import torch, os
+import pandas as pd
+from PIL import Image
+from torchvision.transforms import v2
+from diffsynth.data.video import crop_and_resize
+
+
+class LoraDataset(torch.utils.data.Dataset):
+    def __init__(self, base_path, metadata_path, steps_per_epoch=1000, loras_per_item=1):
+        self.base_path = base_path
+        data_df = pd.read_csv(metadata_path)
+        self.model_file = data_df["model_file"].tolist()
+        self.image_file = data_df["image_file"].tolist()
+        self.text = data_df["text"].tolist()
+        self.max_resolution = 1920 * 1080
+        self.steps_per_epoch = steps_per_epoch
+        self.loras_per_item = loras_per_item
+        
+        
+    def read_image(self, image_file):
+        image = Image.open(image_file).convert("RGB")
+        width, height = image.size
+        if width * height > self.max_resolution:
+            scale = (width * height / self.max_resolution) ** 0.5
+            image = image.resize((int(width / scale), int(height / scale)))
+            width, height = image.size
+        if width % 16 != 0 or height % 16 != 0:
+            image = crop_and_resize(image, height // 16 * 16, width // 16 * 16)
+        image = v2.functional.to_image(image)
+        image = v2.functional.to_dtype(image, dtype=torch.float32, scale=True)
+        image = v2.functional.normalize(image, [0.5], [0.5])
+        return image
+    
+    
+    def get_data(self, data_id):
+        data = {
+            "model_file": os.path.join(self.base_path, self.model_file[data_id]),
+            "image": self.read_image(os.path.join(self.base_path, self.image_file[data_id])),
+            "text": self.text[data_id]
+        }
+        return data
+
+
+    def __getitem__(self, index):
+        data = []
+        while len(data) < self.loras_per_item:
+            data_id = torch.randint(0, len(self.model_file), (1,))[0]
+            data_id = (data_id + index) % len(self.model_file) # For fixed seed.
+            data.append(self.get_data(data_id))
+        return data
+
+
+    def __len__(self):
+        return self.steps_per_epoch

lora/merger.py

@@ -0,0 +1,61 @@
+import torch
+
+
+class LoraMerger(torch.nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight_base = torch.nn.Parameter(torch.randn((dim,)))
+        self.weight_lora = torch.nn.Parameter(torch.randn((dim,)))
+        self.weight_cross = torch.nn.Parameter(torch.randn((dim,)))
+        self.weight_out = torch.nn.Parameter(torch.ones((dim,)))
+        self.bias = torch.nn.Parameter(torch.randn((dim,)))
+        self.activation = torch.nn.Sigmoid()
+        self.norm_base = torch.nn.LayerNorm(dim, eps=1e-5)
+        self.norm_lora = torch.nn.LayerNorm(dim, eps=1e-5)
+        
+    def forward(self, base_output, lora_outputs):
+        norm_base_output = self.norm_base(base_output)
+        norm_lora_outputs = self.norm_lora(lora_outputs)
+        gate = self.activation(
+            norm_base_output * self.weight_base \
+            + norm_lora_outputs * self.weight_lora \
+            + norm_base_output * norm_lora_outputs * self.weight_cross + self.bias
+        )
+        output = base_output + (self.weight_out * gate * lora_outputs).sum(dim=0)
+        return output
+
+
+class LoraPatcher(torch.nn.Module):
+    def __init__(self, lora_patterns=None):
+        super().__init__()
+        if lora_patterns is None:
+            lora_patterns = self.default_lora_patterns()
+        model_dict = {}
+        for lora_pattern in lora_patterns:
+            name, dim = lora_pattern["name"], lora_pattern["dim"]
+            model_dict[name.replace(".", "___")] = LoraMerger(dim)
+        self.model_dict = torch.nn.ModuleDict(model_dict)
+        
+    def default_lora_patterns(self):
+        lora_patterns = []
+        lora_dict = {
+            "attn.a_to_qkv": 9216, "attn.a_to_out": 3072, "ff_a.0": 12288, "ff_a.2": 3072, "norm1_a.linear": 18432,
+            "attn.b_to_qkv": 9216, "attn.b_to_out": 3072, "ff_b.0": 12288, "ff_b.2": 3072, "norm1_b.linear": 18432,
+        }
+        for i in range(19):
+            for suffix in lora_dict:
+                lora_patterns.append({
+                    "name": f"blocks.{i}.{suffix}",
+                    "dim": lora_dict[suffix]
+                })
+        lora_dict = {"to_qkv_mlp": 21504, "proj_out": 3072, "norm.linear": 9216}
+        for i in range(38):
+            for suffix in lora_dict:
+                lora_patterns.append({
+                    "name": f"single_blocks.{i}.{suffix}",
+                    "dim": lora_dict[suffix]
+                })
+        return lora_patterns
+        
+    def forward(self, base_output, lora_outputs, name):
+        return self.model_dict[name.replace(".", "___")](base_output, lora_outputs)

lora/retriever.py

@@ -0,0 +1,149 @@
+import torch
+from diffsynth import SDTextEncoder
+from diffsynth.models.sd3_text_encoder import SD3TextEncoder1StateDictConverter
+from diffsynth.models.sd_text_encoder import CLIPEncoderLayer
+
+
+class LoRALayerBlock(torch.nn.Module):
+    def __init__(self, L, dim_in):
+        super().__init__()
+        self.x = torch.nn.Parameter(torch.randn(1, L, dim_in))
+
+    def forward(self, lora_A, lora_B):
+        out = self.x @ lora_A.T @ lora_B.T
+        return out
+    
+
+class LoRAEmbedder(torch.nn.Module):
+    def __init__(self, lora_patterns=None, L=1, out_dim=2048):
+        super().__init__()
+        if lora_patterns is None:
+            lora_patterns = self.default_lora_patterns()
+            
+        model_dict = {}
+        for lora_pattern in lora_patterns:
+            name, dim = lora_pattern["name"], lora_pattern["dim"][0]
+            model_dict[name.replace(".", "___")] = LoRALayerBlock(L, dim)
+        self.model_dict = torch.nn.ModuleDict(model_dict)
+        
+        proj_dict = {}
+        for lora_pattern in lora_patterns:
+            layer_type, dim = lora_pattern["type"], lora_pattern["dim"][1]
+            if layer_type not in proj_dict:
+                proj_dict[layer_type.replace(".", "___")] = torch.nn.Linear(dim, out_dim)
+        self.proj_dict = torch.nn.ModuleDict(proj_dict)
+        
+        self.lora_patterns = lora_patterns
+        
+        
+    def default_lora_patterns(self):
+        lora_patterns = []
+        lora_dict = {
+            "attn.a_to_qkv": (3072, 9216), "attn.a_to_out": (3072, 3072), "ff_a.0": (3072, 12288), "ff_a.2": (12288, 3072), "norm1_a.linear": (3072, 18432),
+            "attn.b_to_qkv": (3072, 9216), "attn.b_to_out": (3072, 3072), "ff_b.0": (3072, 12288), "ff_b.2": (12288, 3072), "norm1_b.linear": (3072, 18432),
+        }
+        for i in range(19):
+            for suffix in lora_dict:
+                lora_patterns.append({
+                    "name": f"blocks.{i}.{suffix}",
+                    "dim": lora_dict[suffix],
+                    "type": suffix,
+                })
+        lora_dict = {"to_qkv_mlp": (3072, 21504), "proj_out": (15360, 3072), "norm.linear": (3072, 9216)}
+        for i in range(38):
+            for suffix in lora_dict:
+                lora_patterns.append({
+                    "name": f"single_blocks.{i}.{suffix}",
+                    "dim": lora_dict[suffix],
+                    "type": suffix,
+                })
+        return lora_patterns
+        
+    def forward(self, lora):
+        lora_emb = []
+        for lora_pattern in self.lora_patterns:
+            name, layer_type = lora_pattern["name"], lora_pattern["type"]
+            lora_A = lora[name + ".lora_A.default.weight"]
+            lora_B = lora[name + ".lora_B.default.weight"]
+            lora_out = self.model_dict[name.replace(".", "___")](lora_A, lora_B)
+            lora_out = self.proj_dict[layer_type.replace(".", "___")](lora_out)
+            lora_emb.append(lora_out)
+        lora_emb = torch.concat(lora_emb, dim=1)
+        return lora_emb
+
+
+class TextEncoder(torch.nn.Module):
+    def __init__(self, embed_dim=768, vocab_size=49408, max_position_embeddings=77, num_encoder_layers=12, encoder_intermediate_size=3072):
+        super().__init__()
+
+        # token_embedding
+        self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim)
+
+        # position_embeds (This is a fixed tensor)
+        self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim))
+
+        # encoders
+        self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)])
+
+        # attn_mask
+        self.attn_mask = self.attention_mask(max_position_embeddings)
+
+        # final_layer_norm
+        self.final_layer_norm = torch.nn.LayerNorm(embed_dim)
+
+    def attention_mask(self, length):
+        mask = torch.empty(length, length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)
+        return mask
+
+    def forward(self, input_ids, clip_skip=1):
+        embeds = self.token_embedding(input_ids) + self.position_embeds
+        attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype)
+        for encoder_id, encoder in enumerate(self.encoders):
+            embeds = encoder(embeds, attn_mask=attn_mask)
+            if encoder_id + clip_skip == len(self.encoders):
+                break
+        embeds = self.final_layer_norm(embeds)
+        pooled_embeds = embeds[torch.arange(embeds.shape[0]), input_ids.to(dtype=torch.int).argmax(dim=-1)]
+        return pooled_embeds
+    
+    @staticmethod
+    def state_dict_converter():
+        return SD3TextEncoder1StateDictConverter()
+    
+    
+class LoRAEncoder(torch.nn.Module):
+    def __init__(self, embed_dim=768, max_position_embeddings=304, num_encoder_layers=2, encoder_intermediate_size=3072, L=1):
+        super().__init__()
+        max_position_embeddings *= L
+        
+        # Embedder
+        self.embedder = LoRAEmbedder(L=L, out_dim=embed_dim)
+
+        # position_embeds (This is a fixed tensor)
+        self.position_embeds = torch.nn.Parameter(torch.zeros(1, max_position_embeddings, embed_dim))
+
+        # encoders
+        self.encoders = torch.nn.ModuleList([CLIPEncoderLayer(embed_dim, encoder_intermediate_size) for _ in range(num_encoder_layers)])
+
+        # attn_mask
+        self.attn_mask = self.attention_mask(max_position_embeddings)
+
+        # final_layer_norm
+        self.final_layer_norm = torch.nn.LayerNorm(embed_dim)
+
+    def attention_mask(self, length):
+        mask = torch.empty(length, length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)
+        return mask
+
+    def forward(self, lora):
+        embeds = self.embedder(lora) + self.position_embeds
+        attn_mask = self.attn_mask.to(device=embeds.device, dtype=embeds.dtype)
+        for encoder_id, encoder in enumerate(self.encoders):
+            embeds = encoder(embeds, attn_mask=attn_mask)
+        embeds = self.final_layer_norm(embeds)
+        embeds = embeds.mean(dim=1)
+        return embeds

lora/test_merger.py

@@ -0,0 +1,46 @@
+from diffsynth import FluxImagePipeline, ModelManager, load_state_dict
+from diffsynth.models.lora import FluxLoRAConverter
+from diffsynth.pipelines.flux_image import lets_dance_flux
+from lora.dataset import LoraDataset
+from lora.merger import LoraPatcher
+from lora.utils import load_lora
+import torch, os
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from tqdm import tqdm
+
+
+model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda", model_id_list=["FLUX.1-dev"])
+pipe = FluxImagePipeline.from_model_manager(model_manager)
+pipe.enable_auto_lora()
+
+lora_patcher = LoraPatcher().to(dtype=torch.bfloat16, device="cuda")
+lora_patcher.load_state_dict(load_state_dict("models/lora_merger/epoch-3.safetensors"))
+
+dataset = LoraDataset("data/lora/models", "data/lora/lora_dataset_1000.csv", steps_per_epoch=800, loras_per_item=4)
+
+for seed in range(100):
+    batch = dataset[0]
+    num_lora = torch.randint(1, len(batch), (1,))[0]
+    lora_state_dicts = [
+        FluxLoRAConverter.align_to_diffsynth_format(load_lora(batch[i]["model_file"], device="cuda")) for i in range(num_lora)
+    ]
+    image = pipe(
+        prompt=batch[0]["text"],
+        seed=seed,
+    )
+    image.save(f"data/lora/lora_outputs/image_{seed}_nolora.jpg")
+    for i in range(num_lora):
+        image = pipe(
+            prompt=batch[0]["text"],
+            lora_state_dicts=[lora_state_dicts[i]], 
+            lora_patcher=lora_patcher,
+            seed=seed,
+        )
+        image.save(f"data/lora/lora_outputs/image_{seed}_{i}.jpg")
+    image = pipe(
+        prompt=batch[0]["text"],
+        lora_state_dicts=lora_state_dicts, 
+        lora_patcher=lora_patcher,
+        seed=seed,
+    )
+    image.save(f"data/lora/lora_outputs/image_{seed}_merger.jpg")

lora/test_retriever.py

@@ -0,0 +1,148 @@
+from diffsynth import FluxImagePipeline, ModelManager, load_state_dict
+from diffsynth.models.lora import FluxLoRAConverter
+from diffsynth.pipelines.flux_image import lets_dance_flux
+from lora.dataset import LoraDataset
+from lora.retriever import TextEncoder, LoRAEncoder
+from lora.merger import LoraPatcher
+from lora.utils import load_lora
+import torch, os
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from tqdm import tqdm
+from transformers import CLIPTokenizer, CLIPModel
+import pandas as pd
+
+
+
+class LoRARetrieverTrainingModel(torch.nn.Module):
+    def __init__(self, pretrained_path):
+        super().__init__()
+        
+        self.text_encoder = TextEncoder().to(torch.bfloat16)
+        state_dict = load_state_dict("models/FLUX/FLUX.1-dev/text_encoder/model.safetensors")
+        self.text_encoder.load_state_dict(TextEncoder.state_dict_converter().from_civitai(state_dict))
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder.eval()
+        
+        self.lora_encoder = LoRAEncoder().to(torch.bfloat16)
+        state_dict = load_state_dict(pretrained_path)
+        self.lora_encoder.load_state_dict(state_dict)
+        
+        self.tokenizer = CLIPTokenizer.from_pretrained("diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_1")
+        
+        
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.torch_dtype = dtype
+        super().to(*args, **kwargs)
+        return self
+
+
+    def forward(self, batch):
+        text = [data["text"] for data in batch]
+        input_ids = self.tokenizer(
+            text,
+            return_tensors="pt",
+            padding="max_length",
+            max_length=77,
+            truncation=True
+        ).input_ids.to(self.device)
+        text_emb = self.text_encoder(input_ids)
+        text_emb = text_emb / text_emb.norm()
+        
+        lora_emb = []
+        for data in batch:
+            lora = FluxLoRAConverter.align_to_diffsynth_format(load_lora(data["model_file"], device=self.device))
+            lora_emb.append(self.lora_encoder(lora))
+        lora_emb = torch.concat(lora_emb)
+        lora_emb = lora_emb / lora_emb.norm()
+        
+        similarity = text_emb @ lora_emb.T
+        print(similarity)
+        loss = -torch.log(torch.softmax(similarity, dim=0).diag()) - torch.log(torch.softmax(similarity, dim=1).diag())
+        loss = 10 * loss.mean()
+        return loss
+    
+    
+    def trainable_modules(self):
+        return self.lora_encoder.parameters()
+    
+    @torch.no_grad()
+    def process_lora_list(self, lora_list):
+        lora_emb = []
+        for lora in tqdm(lora_list):
+            lora = FluxLoRAConverter.align_to_diffsynth_format(load_lora(lora, device="cuda"))
+            lora_emb.append(self.lora_encoder(lora))
+        lora_emb = torch.concat(lora_emb)
+        lora_emb = lora_emb / lora_emb.norm()
+        self.lora_emb = lora_emb
+        self.lora_list = lora_list
+    
+    @torch.no_grad()
+    def retrieve(self, text, k=1):
+        input_ids = self.tokenizer(
+            text,
+            return_tensors="pt",
+            padding="max_length",
+            max_length=77,
+            truncation=True
+        ).input_ids.to(self.device)
+        text_emb = self.text_encoder(input_ids)
+        text_emb = text_emb / text_emb.norm()
+        
+        similarity = text_emb @ self.lora_emb.T
+        topk = torch.topk(similarity, k, dim=1).indices[0]
+        
+        lora_list = []
+        model_url_list = []
+        for lora_id in topk:
+            print(self.lora_list[lora_id])
+            lora = FluxLoRAConverter.align_to_diffsynth_format(load_lora(self.lora_list[lora_id], device="cuda"))
+            lora_list.append(lora)
+            model_id = self.lora_list[lora_id].split("/")[3:5]
+            model_url_list.append(f"https://www.modelscope.cn/models/{model_id[0]}/{model_id[1]}")
+        return lora_list, model_url_list
+
+
+
+model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda", model_id_list=["FLUX.1-dev"])
+pipe = FluxImagePipeline.from_model_manager(model_manager)
+pipe.enable_auto_lora()
+
+lora_patcher = LoraPatcher().to(dtype=torch.bfloat16, device="cuda")
+lora_patcher.load_state_dict(load_state_dict("models/lora_merger/epoch-9.safetensors"))
+
+retriever = LoRARetrieverTrainingModel("models/lora_retriever/epoch-3.safetensors").to(dtype=torch.bfloat16, device="cuda")
+retriever.process_lora_list(list(set("data/lora/models/" + i for i in pd.read_csv("data/lora/lora_dataset_1000.csv")["model_file"])))
+
+dataset = LoraDataset("data/lora/models", "data/lora/lora_dataset_1000.csv", steps_per_epoch=800, loras_per_item=1)
+
+text_list = []
+model_url_list = []
+for seed in range(100):
+    text = dataset[0][0]["text"]
+    print(text)
+    loras, urls = retriever.retrieve(text, k=3)
+    print(urls)
+    image = pipe(
+        prompt=text,
+        seed=seed,
+    )
+    image.save(f"data/lora/lora_outputs/image_{seed}_top0.jpg")
+    for i in range(2, 3):
+        image = pipe(
+            prompt=text,
+            lora_state_dicts=loras[:i+1],
+            lora_patcher=lora_patcher,
+            seed=seed,
+        )
+        image.save(f"data/lora/lora_outputs/image_{seed}_top{i+1}.jpg")
+        
+    text_list.append(text)
+    model_url_list.append(urls)
+    df = pd.DataFrame()
+    df["text"] = text_list
+    df["models"] = [",".join(i) for i in model_url_list]
+    df.to_csv("data/lora/lora_outputs/metadata.csv", index=False, encoding="utf-8-sig")

lora/train_merger.py

@@ -0,0 +1,119 @@
+from diffsynth import FluxImagePipeline, ModelManager
+from diffsynth.models.lora import FluxLoRAConverter
+from diffsynth.pipelines.flux_image import lets_dance_flux
+from lora.dataset import LoraDataset
+from lora.merger import LoraPatcher
+from lora.utils import load_lora
+import torch, os
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from tqdm import tqdm
+
+
+
+class LoRAMergerTrainingModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu", model_id_list=["FLUX.1-dev"])
+        self.pipe = FluxImagePipeline.from_model_manager(model_manager)
+        self.lora_patcher = LoraPatcher()
+        self.pipe.enable_auto_lora()
+        self.freeze_parameters()
+        self.switch_to_training_mode()
+        self.use_gradient_checkpointing = True
+        self.state_dict_converter = FluxLoRAConverter.align_to_diffsynth_format
+        self.device = "cuda"
+        
+        
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.torch_dtype = dtype
+        super().to(*args, **kwargs)
+        return self
+        
+        
+    def switch_to_training_mode(self):
+        self.pipe.scheduler.set_timesteps(1000, training=True)
+
+
+    def freeze_parameters(self):
+        self.pipe.requires_grad_(False)
+        self.pipe.eval()
+        self.pipe.denoising_model().train()
+        self.lora_patcher.requires_grad_(True)
+
+
+    def forward(self, batch):
+        # Data
+        text, image = batch[0]["text"], batch[0]["image"].unsqueeze(0)
+        num_lora = torch.randint(1, len(batch), (1,))[0]
+        lora_state_dicts = [
+            self.state_dict_converter(load_lora(batch[i]["model_file"], device=self.device)) for i in range(num_lora)
+        ]
+        lora_alphas = None
+
+        # Prepare input parameters
+        self.pipe.device = self.device
+        prompt_emb = self.pipe.encode_prompt(text, positive=True)
+        latents = self.pipe.vae_encoder(image.to(dtype=self.pipe.torch_dtype, device=self.device))
+        noise = torch.randn_like(latents)
+        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))
+        timestep = self.pipe.scheduler.timesteps[timestep_id].to(self.device)
+        extra_input = self.pipe.prepare_extra_input(latents)
+        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)
+        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)
+
+        # Compute loss
+        noise_pred = lets_dance_flux(
+            self.pipe.dit,
+            hidden_states=noisy_latents, timestep=timestep, **prompt_emb, **extra_input,
+            lora_state_dicts=lora_state_dicts, lora_alphas=lora_alphas, lora_patcher=self.lora_patcher,
+            use_gradient_checkpointing=self.use_gradient_checkpointing
+        )
+        loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())
+        loss = loss * self.pipe.scheduler.training_weight(timestep)
+        return loss
+    
+    
+    def trainable_modules(self):
+        return self.lora_patcher.parameters()
+
+
+class ModelLogger:
+    def __init__(self, output_path, remove_prefix_in_ckpt=None):
+        self.output_path = output_path
+        self.remove_prefix_in_ckpt = remove_prefix_in_ckpt
+        
+    
+    def on_step_end(self, loss):
+        pass
+    
+    
+    def on_epoch_end(self, accelerator, model, epoch_id):
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            state_dict = accelerator.unwrap_model(model).lora_patcher.state_dict()
+            os.makedirs(self.output_path, exist_ok=True)
+            path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors")
+            accelerator.save(state_dict, path, safe_serialization=True)
+
+
+if __name__ == '__main__':
+    model = LoRAMergerTrainingModel()
+    dataset = LoraDataset("data/lora/models/", "data/lora/lora_dataset_1000.csv", steps_per_epoch=800, loras_per_item=4)
+    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=1, num_workers=1, collate_fn=lambda x: x[0])
+    optimizer = torch.optim.AdamW(model.trainable_modules(), lr=1e-4)
+    model_logger = ModelLogger("models/lora_merger")
+    accelerator = Accelerator(kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)])
+    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+    
+    for epoch_id in range(1000000):
+        for data in tqdm(dataloader):
+            with accelerator.accumulate(model):
+                optimizer.zero_grad()
+                loss = model(data)
+                accelerator.backward(loss)
+                optimizer.step()
+        model_logger.on_epoch_end(accelerator, model, epoch_id)

lora/train_retriever.py

@@ -0,0 +1,105 @@
+from diffsynth import FluxImagePipeline, ModelManager, load_state_dict
+from diffsynth.models.lora import FluxLoRAConverter
+from diffsynth.pipelines.flux_image import lets_dance_flux
+from lora.dataset import LoraDataset
+from lora.retriever import TextEncoder, LoRAEncoder
+from lora.utils import load_lora
+import torch, os
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from tqdm import tqdm
+from transformers import CLIPTokenizer, CLIPModel
+
+
+
+class LoRARetrieverTrainingModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        
+        self.text_encoder = TextEncoder().to(torch.bfloat16)
+        state_dict = load_state_dict("models/FLUX/FLUX.1-dev/text_encoder/model.safetensors")
+        self.text_encoder.load_state_dict(TextEncoder.state_dict_converter().from_civitai(state_dict))
+        self.text_encoder.requires_grad_(False)
+        self.text_encoder.eval()
+        
+        self.lora_encoder = LoRAEncoder().to(torch.bfloat16)
+        
+        self.tokenizer = CLIPTokenizer.from_pretrained("diffsynth/tokenizer_configs/stable_diffusion_3/tokenizer_1")
+        
+        
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
+        if device is not None:
+            self.device = device
+        if dtype is not None:
+            self.torch_dtype = dtype
+        super().to(*args, **kwargs)
+        return self
+
+
+    def forward(self, batch):
+        text = [data["text"] for data in batch]
+        input_ids = self.tokenizer(
+            text,
+            return_tensors="pt",
+            padding="max_length",
+            max_length=77,
+            truncation=True
+        ).input_ids.to(self.device)
+        text_emb = self.text_encoder(input_ids)
+        text_emb = text_emb / text_emb.norm()
+        
+        lora_emb = []
+        for data in batch:
+            lora = FluxLoRAConverter.align_to_diffsynth_format(load_lora(data["model_file"], device=self.device))
+            lora_emb.append(self.lora_encoder(lora))
+        lora_emb = torch.concat(lora_emb)
+        lora_emb = lora_emb / lora_emb.norm()
+        
+        similarity = text_emb @ lora_emb.T
+        print(similarity)
+        loss = -torch.log(torch.softmax(similarity, dim=0).diag()) - torch.log(torch.softmax(similarity, dim=1).diag())
+        loss = 10 * loss.mean()
+        return loss
+    
+    
+    def trainable_modules(self):
+        return self.lora_encoder.parameters()
+
+
+class ModelLogger:
+    def __init__(self, output_path, remove_prefix_in_ckpt=None):
+        self.output_path = output_path
+        self.remove_prefix_in_ckpt = remove_prefix_in_ckpt
+        
+    
+    def on_step_end(self, loss):
+        pass
+    
+    
+    def on_epoch_end(self, accelerator, model, epoch_id):
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            state_dict = accelerator.unwrap_model(model).lora_encoder.state_dict()
+            os.makedirs(self.output_path, exist_ok=True)
+            path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors")
+            accelerator.save(state_dict, path, safe_serialization=True)
+
+
+if __name__ == '__main__':
+    model = LoRARetrieverTrainingModel()
+    dataset = LoraDataset("data/lora/models/", "data/lora/lora_dataset_1000.csv", steps_per_epoch=100, loras_per_item=32)
+    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=1, num_workers=1, collate_fn=lambda x: x[0])
+    optimizer = torch.optim.AdamW(model.trainable_modules(), lr=1e-4)
+    model_logger = ModelLogger("models/lora_retriever")
+    accelerator = Accelerator(kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)])
+    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
+    
+    for epoch_id in range(1000000):
+        for data in tqdm(dataloader):
+            with accelerator.accumulate(model):
+                optimizer.zero_grad()
+                loss = model(data)
+                accelerator.backward(loss)
+                optimizer.step()
+                print(loss)
+        model_logger.on_epoch_end(accelerator, model, epoch_id)

lora/utils.py

@@ -0,0 +1,12 @@
+from diffsynth import load_state_dict
+import math, torch
+
+
+def load_lora(file_path, device):
+    sd = load_state_dict(file_path, torch_dtype=torch.bfloat16, device=device)
+    scale = math.sqrt(sd["lora_unet_single_blocks_9_modulation_lin.alpha"] / sd["lora_unet_single_blocks_9_modulation_lin.lora_down.weight"].shape[0])
+    if scale != 1:
+        sd = {i: sd[i] * scale for i in sd}
+    return sd
+
+