mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-13 23:53:37 +00:00
Need to test saving and loading.
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user