Need to test saving and loading.

This commit is contained in:
CrossProduct
2025-01-18 06:58:04 +00:00
parent 3c605a124e
commit 48d2b61029
@@ -56,6 +56,27 @@ class LTXVideoConditionedTransformer3DModel(LTXVideoTransformer3DModel):
attention_out_bias=attention_out_bias
)
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
model = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
# model is now a LTXVideoTransformer3DModel instance;
# but we want an instance of LTXVideoTransformer3DModelWithAdapter
# so we often do something like:
new_model = cls(**model.config.__dict__)
new_model.load_state_dict(model.state_dict())
return new_model
def save_pretrained(self, save_directory: str, **kwargs):
"""
Saves model weights (including adapter) + config to disk
in a format compatible with .from_pretrained()
"""
# 1) Save config
self.config.save_pretrained(save_directory)
# 2) Save PyTorch state dict
state_dict = self.state_dict()
torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))
def forward(
self,