Adapter weights can be loaded.

finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py

@@ -1,6 +1,6 @@
 import math
 from typing import Any, Dict, Optional, Tuple
-
+import os
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -57,8 +57,8 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
             adapter_init_scale=1e-3
         )
 
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+    @classmethod # rewrite this to ensure that if it doesn't have adapter weights it will load this way but if it has adapter weights it will not.
+    def from_pretrained(cls, pretrained_model_name_or_path,save_directory, **kwargs):
         # First create an empty model with the desired architecture
         print("Pretrain Loading")
         model = cls()
@@ -73,12 +73,34 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
         filtered_dict = {}
         for k, v in pretrained_dict.items():
             if k in model_dict:
+                
                 filtered_dict[k] = v
+            else:
+                print(k)
         # Update model with pretrained weights
         model_dict.update(filtered_dict)
+
+        adapter_weights_path = os.path.join(save_directory, "adapter_weights.pth")
+        if os.path.exists(adapter_weights_path):
+            adapter_weights = torch.load(adapter_weights_path)
+            model.adapter.load_state_dict(adapter_weights)
+            print("Adapter weights loaded successfully.")
+        else:
+            print("No adapter weights file found. Using default adapter weights.")
+
         model.load_state_dict(model_dict)
         return model
 
+    def save_pretrained(self,save_directory):
+        model_state_dict = self.state_dict()
+        torch.save(model_state_dict, save_directory)
+        
+        adapter_state_dict = self.adapter.state_dict()
+        save_directory = os.path.dirname(save_directory)
+        path = os.path.join(save_directory, "adapter_weights.pth")
+        torch.save(adapter_state_dict, path)
+        print("Saving Model with Adapter")
+
     def forward(
             self,
             hidden_states: torch.Tensor,