CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-08-31 00:10:39 +02:00
Code Issues Packages Projects Releases Wiki Activity

Always config injection

Browse Source
This commit is contained in:
henryruhs
2025-03-11 14:43:10 +01:00
parent 368da824aa
commit b59e172fa3
1 changed files with 17 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
embedding_converter/src/training.py
+17 -18
View File
@@ -2,9 +2,8 @@ import configparser
import os import os
from typing import Tuple from typing import Tuple
import lightning
import torch import torch
from lightning import Trainer from lightning import Trainer, LightningModule
from lightning.pytorch.callbacks import ModelCheckpoint from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import TensorBoardLogger from lightning.pytorch.loggers import TensorBoardLogger
from torch import Tensor, nn from torch import Tensor, nn
@@ -13,22 +12,16 @@ from torchdata.stateful_dataloader import StatefulDataLoader
from .dataset import StaticDataset from .dataset import StaticDataset
from .models.embedding_converter import EmbeddingConverter from .models.embedding_converter import EmbeddingConverter
from .types import Batch, Embedding, OptimizerSet from .types import Batch, Embedding, OptimizerSet, Config
CONFIG = configparser.ConfigParser() CONFIG = configparser.ConfigParser()
CONFIG.read('config.ini') CONFIG.read('config.ini')
class EmbeddingConverterTrainer(lightning.LightningModule): class EmbeddingConverterTrainer(LightningModule):
def __init__(self) -> None: def __init__(self, config : Config) -> None:
super(EmbeddingConverterTrainer, self).__init__() super().__init__()
self.config =\ self.config = config
{
'source_path': CONFIG.get('training.model', 'source_path'),
'target_path': CONFIG.get('training.model', 'target_path'),
'learning_rate': CONFIG.getfloat('training.trainer', 'learning_rate')
}
self.embedding_converter = EmbeddingConverter() self.embedding_converter = EmbeddingConverter()
self.source_embedder = torch.jit.load(self.config.get('source_path'), map_location = 'cpu') # type:ignore[no-untyped-call] self.source_embedder = torch.jit.load(self.config.get('source_path'), map_location = 'cpu') # type:ignore[no-untyped-call]
self.target_embedder = torch.jit.load(self.config.get('target_path'), map_location = 'cpu') # type:ignore[no-untyped-call] self.target_embedder = torch.jit.load(self.config.get('target_path'), map_location = 'cpu') # type:ignore[no-untyped-call]
@@ -128,21 +121,27 @@ def create_trainer() -> Trainer:
def train() -> None: def train() -> None:
config =\ config_dataset =\
{ {
'file_pattern': CONFIG.get('training.dataset', 'file_pattern'), 'file_pattern': CONFIG.get('training.dataset', 'file_pattern'),
'resume_path': CONFIG.get('training.output', 'resume_path') 'resume_path': CONFIG.get('training.output', 'resume_path')
} }
config_trainer =\
{
'source_path': CONFIG.get('training.model', 'source_path'),
'target_path': CONFIG.get('training.model', 'target_path'),
'learning_rate': CONFIG.getfloat('training.trainer', 'learning_rate')
}
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.set_float32_matmul_precision('high') torch.set_float32_matmul_precision('high')
dataset = StaticDataset(config) dataset = StaticDataset(config_dataset)
training_loader, validation_loader = create_loaders(dataset) training_loader, validation_loader = create_loaders(dataset)
embedding_converter_trainer = EmbeddingConverterTrainer() embedding_converter_trainer = EmbeddingConverterTrainer(config_trainer)
trainer = create_trainer() trainer = create_trainer()
if os.path.exists(config.get('resume_path')): if os.path.exists(config_dataset.get('resume_path')):
trainer.fit(embedding_converter_trainer, training_loader, validation_loader, ckpt_path = config.get('resume_path')) trainer.fit(embedding_converter_trainer, training_loader, validation_loader, ckpt_path = config_dataset.get('resume_path'))
else: else:
trainer.fit(embedding_converter_trainer, training_loader, validation_loader) trainer.fit(embedding_converter_trainer, training_loader, validation_loader)