diff --git a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py index 2fcce45..f202e07 100644 --- a/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py +++ b/finetrainers/conditioning/LTXVideoConditionedTransformer3DModel.py @@ -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,