Merge pull request #10 from neph1/update-v0.8.0

Update v0.8.0

Author: neph1

app.py

@@ -17,7 +17,7 @@ class App:
 
     def __init__(self):
         self.configs_path = "config/"
-        self.tabs = dict() # Type Tab
+        self.tabs = dict()
         self.setup_views()
 
     def setup_views(self):

editor_factory.py

@@ -0,0 +1,11 @@
+[project]
+name = "finetrainers-ui"
+version = "0.8.0"
+dependencies = [
+    "gradio",
+    "torch>=2.4.1"
+]
+
+
+[project.urls]
+Homepage = "https://github.com/neph1/finetrainers-ui"

gradio_functions.py

@@ -94,10 +94,15 @@ def run(self, config: Config, finetrainers_path: str, log_file: str):
         return "Unknown result"
 
     def stop(self):
-        self.running = False
-        if self.process:
-            self.process.terminate()
-            time.sleep(3)
-            if self.process.poll() is None:
-                self.process.kill()
-        return "Training stopped"
+        try:
+            self.running = False
+            if self.process:
+                self.process.terminate()
+                time.sleep(3)
+                if self.process.poll() is None:
+                    self.process.kill()
+        except Exception as e:
+            return f"Error stopping training: {e}"
+        finally:
+            self.process.wait()
+        return "Training forcibly stopped"

pyproject.toml

@@ -2,7 +2,6 @@
 from typing import OrderedDict
 import gradio as gr
 import yaml
-import editor_factory
 
 class Tab(ABC):
 

run_trainer.py

@@ -3,11 +3,12 @@
 import os
 import gradio as gr
 from typing import OrderedDict
-from config import Config, global_config
+from config import Config
 
 from run_trainer import RunTrainer
 from tabs import general_tab
 from tabs.tab import Tab
+from trainer_config_validator import TrainerValidator
 
 properties = OrderedDict()
 
@@ -69,16 +70,21 @@ def run_trainer(self, *args):
             key = keys_list[index]
             properties[key].value = properties_values[index]
             config.set(key, properties_values[index])
+        config.set('path_to_finetrainers', general_tab.properties['path_to_finetrainers'].value)
+
+        config_validator = TrainerValidator(config)
+        try:
+            config_validator.validate()
+        except Exception as e:
+            return str(e), None
 
         output_path = os.path.join(properties['output_dir'].value, "config")
         os.makedirs(output_path, exist_ok=True)
         self.save_edits(os.path.join(output_path, "config_{}.yaml".format(time)), *properties_values)
 
         log_file = os.path.join(output_path, "log_{}.txt".format(time))
 
-        if not general_tab.properties['path_to_finetrainers'].value:
-            return "Please set the path to finetrainers in General Settings"
-        result = self.trainer.run(config, general_tab.properties['path_to_finetrainers'].value, log_file)
+        result = self.trainer.run(config, config.get('path_to_finetrainers'), log_file)
         self.trainer.running = False
         if isinstance(result, str):
             return result, log_file

tabs/tab.py

@@ -0,0 +1,78 @@
+import os
+import pytest
+from unittest.mock import patch
+
+from trainer_config_validator import TrainerValidator
+
+@pytest.fixture
+def valid_config():
+    return {
+        'path_to_finetrainers': '/path/to/finetrainers',
+        'accelerate_config': 'config1',
+        'batch_size': 32,
+        'beta1': 0.9,
+        'beta2': 0.999,
+        'caption_column': 'captions.txt',
+        'caption_dropout_p': 0.1,
+        'checkpointing_limit': 5,
+        'checkpointing_steps': 1000,
+        'data_root': '/path/to/data',
+        'dataloader_num_workers': 0,
+        'epsilon': 1e-8,
+        'gpu_ids': '0,1',
+        'gradient_accumulation_steps': 2,
+        'gradient_checkpointing': True,
+        'id_token': 'token123',
+        'lora_alpha': 128,
+        'lr': 0.001,
+        'lr_num_cycles': 10,
+        'lr_scheduler': 'scheduler1',
+        'lr_warmup_steps': 500,
+        'max_grad_norm': 1.0,
+        'mixed_precision': 'fp16',
+        'model_name': 'model_v1',
+        'nccl_timeout': 60,
+        'optimizer': 'adam',
+        'pretrained_model_name_or_path': 'pretrained_model',
+        'rank': 64,
+        'seed': 42,
+        'target_modules': 'module1',
+        'tracker_name': 'tracker',
+        'train_steps': 10000,
+        'training_type': 'type1',
+        'validation_steps': 100,
+        'video_column': 'videos.txt',
+        'video_resolution_buckets': '24x480x720',
+        'weight_decay': 0.01
+    }
+
+@pytest.fixture
+def trainer_validator(valid_config):
+    return TrainerValidator(valid_config)
+
+def test_valid_config(valid_config):
+    trainer_validator = TrainerValidator(valid_config)
+    with patch('os.path.isfile', return_value=True), patch('os.path.exists', return_value=True), patch('os.path.isdir', return_value=True):
+        trainer_validator.validate()
+
+def test_validate_data_root_invalid(trainer_validator):
+    trainer_validator.config['data_root'] = '/invalid/path'
+    with pytest.raises(ValueError, match="data_root path /invalid/path does not exist"):
+        trainer_validator.validate_data_root()
+
+def test_validate_data_root_valid(trainer_validator):
+    with patch('os.path.exists', return_value=True), patch('os.path.isdir', return_value=True):
+        trainer_validator.config['data_root'] = '/path/to/data'
+        trainer_validator.validate_data_root()
+
+def test_validate_video_resolution_buckets_invalid(trainer_validator):
+    trainer_validator.config['video_resolution_buckets'] = '720p,1080p,4k'
+    with pytest.raises(ValueError, match=f"Each bucket must have the format '<frames>x<height>x<width>', but got {trainer_validator.config['video_resolution_buckets']}"):
+        trainer_validator.validate_video_resolution_buckets()
+
+def test_validate_video_resolution_buckets_valid(trainer_validator):
+    trainer_validator.config['video_resolution_buckets'] = '24x480x720'
+    trainer_validator.validate_video_resolution_buckets()
+
+    trainer_validator.config['video_resolution_buckets'] = '8x320x512 24x480x720 30x720x1280'
+    trainer_validator.validate_video_resolution_buckets()

tabs/training_tab.py

@@ -0,0 +1,89 @@
+import os
+import re
+
+class TrainerValidator:
+
+    def __init__(self, config):
+        self.config = config
+
+    def validate(self):
+        required_settings = [
+            'path_to_finetrainers',
+            'accelerate_config',
+            'batch_size',
+            'beta1',
+            'beta2',
+            'caption_column',
+            'caption_dropout_p',
+            'checkpointing_limit',
+            'checkpointing_steps',
+            'data_root',
+            'dataloader_num_workers',
+            'epsilon',
+            'gpu_ids',
+            'gradient_accumulation_steps',
+            'id_token',
+            'lora_alpha',
+            'lr',
+            'lr_num_cycles',
+            'lr_scheduler',
+            'lr_warmup_steps',
+            'max_grad_norm',
+            'mixed_precision',
+            'model_name',
+            'nccl_timeout',
+            'optimizer',
+            'pretrained_model_name_or_path',
+            'rank',
+            'seed',
+            'target_modules',
+            'train_steps',
+            'training_type',
+            'validation_steps',
+            'video_column',
+            'video_resolution_buckets',
+            'weight_decay'
+        ]
+
+        for setting in required_settings:
+            if not self.config.get(setting) and self.config.get(setting) != 0:
+                raise ValueError(f"{setting} is required")
+            
+        self.validate_finetrainers_path()
+        self.validate_data_root()
+        self.validate_caption_column()
+        self.validate_video_column()
+        self.validate_video_resolution_buckets()
+
+    def validate_finetrainers_path(self):
+        train_script_path = os.path.join(self.config.get('path_to_finetrainers'), 'train.py')
+        if not os.path.isfile(train_script_path):
+            raise ValueError(f"train.py does not exist at {self.config.get('path_to_finetrainers')}")
+
+    def validate_caption_column(self):
+        data_root = self.config.get('data_root')
+        caption_column = self.config.get('caption_column')
+        if data_root and caption_column:
+            file_path = os.path.join(data_root, caption_column)
+            if not os.path.isfile(file_path):
+                raise ValueError(f"File {caption_column} does not exist at {data_root}")
+
+    def validate_data_root(self):
+        data_root = self.config.get('data_root')
+        if data_root and not os.path.isdir(data_root):
+            raise ValueError(f"data_root path {data_root} does not exist")
+
+    def validate_video_column(self):
+        data_root = self.config.get('data_root')
+        video_column = self.config.get('video_column')
+        if data_root and video_column:
+            file_path = os.path.join(data_root, video_column)
+            if not os.path.isfile(file_path):
+                raise ValueError(f"File {video_column} does not exist at {data_root}")
+
+    def validate_video_resolution_buckets(self):
+        buckets = self.config.get('video_resolution_buckets')
+        split_buckets =  buckets.split(' ')
+        for bucket in split_buckets:
+            if not isinstance(bucket, str) or not re.match(r'^\d+x\d+x\d+$', bucket):
+                raise ValueError(f"Each bucket must have the format '<frames>x<height>x<width>', but got {bucket}")