Merge branch 'main' into update-v0.10.0

Author: neph1

.github/workflows/python-app.yml

@@ -0,0 +1,39 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Python application
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install flake8 pytest
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Lint with flake8
+      run: |
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - name: Test with pytest
+      run: |
+        pytest

README.md

@@ -12,6 +12,8 @@ I try to keep main stable, but if it fails, step back one version and try that.
 
 ![Screenshot from 2024-12-30 07-56-37](https://github.com/user-attachments/assets/91b947db-1e50-42e0-8d12-28b436bf837d)
 
+v0.9.1: Add 'resume_from_checkpoint' setting
+
 v0.9.0: Add missing precompute_condition settings
 
 v0.8.0: Configuration validator. Fail early.

config/config_categories.yaml

@@ -1,5 +1,5 @@
-Dataset: data_root, video_column, caption_column, id_token, video_resolution_buckets, caption_dropout_p, precompute_conditions
-Training: training_type, seed, train_steps, rank, lora_alpha, target_modules, gradient_accumulation_steps, checkpointing_steps, checkpointing_limit, enable_slicing, enable_tiling, batch_size
+Dataset: data_root, video_column, caption_column, dataset_file, id_token, image_resolution_buckets, video_resolution_buckets, caption_dropout_p, precompute_conditions
+Training: training_type, seed, train_steps, rank, lora_alpha, target_modules, gradient_accumulation_steps, checkpointing_steps, checkpointing_limit, enable_slicing, enable_tiling, batch_size, resume_from_checkpoint
 Optimizer: optimizer, lr, beta1, beta2, epsilon, weight_decay, max_grad_norm, lr_scheduler, lr_num_cycles, lr_warmup_steps
 Validation: validation_steps, validation_epochs, num_validation_videos, validation_prompts, validation_prompt_separator
 Accelerate: gpu_ids, nccl_timeout, gradient_checkpointing, allow_tf32, dataloader_num_workers, report_to, accelerate_config

config/config_template.yaml

@@ -10,7 +10,9 @@ checkpointing_limit: 102
 checkpointing_steps: 500
 data_root: ''
 dataloader_num_workers: 0
+dataset_file: ''
 diffusion_options: ''
+enable_model_cpu_offload: false
 enable_slicing: true
 enable_tiling: true
 epsilon: 1e-8
@@ -21,6 +23,7 @@ id_token: afkx
 layerwise_upcasting_modules: [none, transformer]
 layerwise_upcasting_granularity: [pytorch_layer, diffusers_layer]
 layerwise_upcasting_storage_dtype: [float8_e4m3fn, float8_e5m2]
+image_resolution_buckets: 512x768
 lora_alpha: 128
 lr: 0.0001
 lr_num_cycles: 1
@@ -37,6 +40,7 @@ precompute_conditions: false
 pretrained_model_name_or_path: ''
 rank: 128
 report_to: none
+resume_from_checkpoint: ''
 seed: 42
 target_modules: to_q to_k to_v to_out.0
 text_encoder_dtype: [bf16, fp16, fp32, fp8]

pyproject.toml

@@ -5,7 +5,13 @@ dependencies = [
     "gradio",
     "torch>=2.4.1"
 ]
+description = "A gradio based ui for training video transformer models with finetrainers as backend"
+readme = "README.md"
+license = {file = "LICENSE"}
 
 
 [project.urls]
-Homepage = "https://github.com/neph1/finetrainers-ui"
+Repository = "https://github.com/neph1/finetrainers-ui"
+
+[tool.setuptools]
+packages = ["tabs", "config"]

run_trainer.py

@@ -1,7 +1,8 @@
 import os
 import signal
 import subprocess
-import time
+
+import psutil
 
 from config import Config
 
@@ -18,88 +19,92 @@ def run(self, config: Config, finetrainers_path: str, log_file: str):
         assert config.get('data_root'), "Data root required"
         assert config.get('pretrained_model_name_or_path'), "pretrained_model_name_or_path required"
 
-        # Model arguments
-        model_cmd = f"--model_name {config.get('model_name')} \
-        --pretrained_model_name_or_path {config.get('pretrained_model_name_or_path')} \
-        --text_encoder_dtype {config.get('text_encoder_dtype')} \
-        --text_encoder_2_dtype {config.get('text_encoder_2_dtype')} \
-        --text_encoder_3_dtype {config.get('text_encoder_3_dtype')} \
-        --vae_dtype {config.get('vae_dtype')} "
-        
+        model_cmd = ["--model_name", config.get('model_name'), 
+                     "--pretrained_model_name_or_path", config.get('pretrained_model_name_or_path'),
+                     "--text_encoder_dtype", config.get('text_encoder_dtype'),
+                     "--text_encoder_2_dtype", config.get('text_encoder_2_dtype'),
+                     "--text_encoder_3_dtype", config.get('text_encoder_3_dtype'),
+                     "--vae_dtype", config.get('vae_dtype')]
+
         if config.get('layerwise_upcasting_modules') != 'none':
-            model_cmd += f"--layerwise_upcasting_modules {config.get('layerwise_upcasting_modules')} \
-            --layerwise_upcasting_storage_dtype {config.get('layerwise_upcasting_storage_dtype')} \
-            --layerwise_upcasting_granularity {config.get('layerwise_upcasting_granularity')} "
-
-        # Dataset arguments
-        dataset_cmd = f"--data_root {config.get('data_root')} \
-        --video_column {config.get('video_column')} \
-        --caption_column {config.get('caption_column')} \
-        --id_token {config.get('id_token')} \
-        --video_resolution_buckets {config.get('video_resolution_buckets')} \
-        --caption_dropout_p {config.get('caption_dropout_p')} \
-        --caption_dropout_technique {config.get('caption_dropout_technique')} \
-        {'--precompute_conditions' if config.get('precompute_conditions') else ''} "
-
-        # Dataloader arguments
-        dataloader_cmd = f"--dataloader_num_workers {config.get('dataloader_num_workers')}"
+            model_cmd +=["--layerwise_upcasting_modules", config.get('layerwise_upcasting_modules'),
+            "--layerwise_upcasting_storage_dtype", config.get('layerwise_upcasting_storage_dtype'),
+            "--layerwise_upcasting_granularity", config.get('layerwise_upcasting_granularity')]
+
+        dataset_cmd = ["--data_root", config.get('data_root'),
+                   "--video_column", config.get('video_column'),
+                   "--caption_column", config.get('caption_column'),
+                   "--id_token", config.get('id_token'),
+                   "--video_resolution_buckets"]
+        dataset_cmd += config.get('video_resolution_buckets').split(' ')
+        dataset_cmd += ["--image_resolution_buckets"]
+        dataset_cmd += config.get('image_resolution_buckets').split(' ')
+        dataset_cmd += ["--caption_dropout_p", config.get('caption_dropout_p'),
+                   "--caption_dropout_technique", config.get('caption_dropout_technique'),
+                   "--text_encoder_dtype", config.get('text_encoder_dtype'),
+                   "--text_encoder_2_dtype", config.get('text_encoder_2_dtype'),
+                   "--text_encoder_3_dtype", config.get('text_encoder_3_dtype'),
+                   "--vae_dtype", config.get('vae_dtype'),
+                   '--precompute_conditions' if config.get('precompute_conditions') else '']
+        if config.get('dataset_file'):
+            dataset_cmd += ["--dataset_file", config.get('dataset_file')]
+
+        dataloader_cmd = ["--dataloader_num_workers", config.get('dataloader_num_workers')]
 
         # Diffusion arguments TODO: replace later
-        diffusion_cmd = f"{config.get('diffusion_options')}"
-
-        # Training arguments
-        training_cmd = f"--training_type {config.get('training_type')} \
-        --seed {config.get('seed')} \
-        --batch_size {config.get('batch_size')} \
-        --train_steps {config.get('train_steps')} \
-        --rank {config.get('rank')} \
-        --lora_alpha {config.get('lora_alpha')} \
-        --target_modules {config.get('target_modules')} \
-        --gradient_accumulation_steps {config.get('gradient_accumulation_steps')} \
-        {'--gradient_checkpointing' if config.get('gradient_checkpointing') else ''} \
-        --checkpointing_steps {config.get('checkpointing_steps')} \
-        --checkpointing_limit {config.get('checkpointing_limit')} \
-        {'--enable_slicing' if config.get('enable_slicing') else ''} \
-        {'--enable_tiling' if config.get('enable_tiling') else ''}"
-
-        # Optimizer arguments
-        optimizer_cmd = f"--optimizer {config.get('optimizer')} \
-        --lr {config.get('lr')} \
-        --lr_scheduler {config.get('lr_scheduler')} \
-        --lr_warmup_steps {config.get('lr_warmup_steps')} \
-        --lr_num_cycles {config.get('lr_num_cycles')} \
-        --beta1 {config.get('beta1')} \
-        --beta2 {config.get('beta2')} \
-        --weight_decay {config.get('weight_decay')} \
-        --epsilon {config.get('epsilon')} \
-        --max_grad_norm {config.get('max_grad_norm')} \
-        {'--use_8bit_bnb' if config.get('use_8bit_bnb') else ''}"
-
-        # Validation arguments
-        validation_cmd = f"--validation_prompts \"{config.get('validation_prompts')}\" \
-        --num_validation_videos {config.get('num_validation_videos')} \
-        --validation_steps {config.get('validation_steps')}"
-
-        # Miscellaneous arguments
-        miscellaneous_cmd = f"--tracker_name {config.get('tracker_name')} \
-        --output_dir {config.get('output_dir')} \
-        --nccl_timeout {config.get('nccl_timeout')} \
-        --report_to {config.get('report_to')}"
-
-        cmd = f"accelerate launch --config_file {finetrainers_path}/accelerate_configs/{config.get('accelerate_config')} --gpu_ids {config.get('gpu_ids')} {finetrainers_path}/train.py \
-        {model_cmd} \
-        {dataset_cmd} \
-        {dataloader_cmd} \
-        {diffusion_cmd} \
-        {training_cmd} \
-        {optimizer_cmd} \
-        {validation_cmd} \
-        {miscellaneous_cmd}"
-
-        print(cmd)
+        diffusion_cmd = [config.get('diffusion_options')]
+
+        training_cmd = ["--training_type", config.get('training_type'),
+                "--seed", config.get('seed'),
+                "--mixed_precision", config.get('mixed_precision'),
+                "--batch_size", config.get('batch_size'),
+                "--train_steps", config.get('train_steps'),
+                "--rank", config.get('rank'),
+                "--lora_alpha", config.get('lora_alpha'),
+                "--target_modules"]
+        training_cmd += config.get('target_modules').split(' ')
+        training_cmd += ["--gradient_accumulation_steps", config.get('gradient_accumulation_steps'),
+                '--gradient_checkpointing' if config.get('gradient_checkpointing') else '',
+                "--checkpointing_steps", config.get('checkpointing_steps'),
+                "--checkpointing_limit", config.get('checkpointing_limit'),
+                '--enable_slicing' if config.get('enable_slicing') else '',
+                '--enable_tiling' if config.get('enable_tiling') else '']
+        if config.get('enable_model_cpu_offload'):
+            training_cmd += ["--enable_model_cpu_offload"]
+
+        if config.get('resume_from_checkpoint'):
+            training_cmd += ["--resume_from_checkpoint", config.get('resume_from_checkpoint')]
+
+        optimizer_cmd = ["--optimizer", config.get('optimizer'),
+                 "--lr", config.get('lr'),
+                 "--lr_scheduler", config.get('lr_scheduler'),
+                 "--lr_warmup_steps", config.get('lr_warmup_steps'),
+                 "--lr_num_cycles", config.get('lr_num_cycles'),
+                 "--beta1", config.get('beta1'),
+                 "--beta2", config.get('beta2'),
+                 "--weight_decay", config.get('weight_decay'),
+                 "--epsilon", config.get('epsilon'),
+                 "--max_grad_norm", config.get('max_grad_norm'),
+                 '--use_8bit_bnb' if config.get('use_8bit_bnb') else '']
+
+        validation_cmd = ["--validation_prompts" if config.get('validation_prompts') else '', config.get('validation_prompts') or '',
+                  "--num_validation_videos", config.get('num_validation_videos'),
+                  "--validation_steps", config.get('validation_steps')]
+
+        miscellaneous_cmd = ["--tracker_name", config.get('tracker_name'),
+                     "--output_dir", config.get('output_dir'),
+                     "--nccl_timeout", config.get('nccl_timeout'),
+                     "--report_to", config.get('report_to')]
+        accelerate_cmd = ["accelerate", "launch", "--config_file", f"{finetrainers_path}/accelerate_configs/{config.get('accelerate_config')}", "--gpu_ids", config.get('gpu_ids')]
+        cmd = accelerate_cmd + [f"{finetrainers_path}/train.py"] + model_cmd + dataset_cmd + dataloader_cmd + diffusion_cmd + training_cmd + optimizer_cmd + validation_cmd + miscellaneous_cmd
+        fixed_cmd = []
+        for i in range(len(cmd)):
+            if cmd[i] != '':
+                fixed_cmd.append(f"{cmd[i]}")
+        print(' '.join(fixed_cmd))
         self.running = True
         with open(log_file, "w") as output_file:
-            self.process = subprocess.Popen(cmd, shell=True, stdout=output_file, stderr=output_file, text=True)
+            self.process = subprocess.Popen(fixed_cmd, shell=False, stdout=output_file, stderr=output_file, text=True, preexec_fn=os.setsid)
             self.process.communicate()
             return self.process
 
@@ -109,12 +114,20 @@ def stop(self):
         try:
             self.running = False
             if self.process:
-                self.process.terminate()
-                time.sleep(3)
-                if self.process.poll() is None:
-                    self.process.kill()
+                os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
+                self.terminate_process_tree(self.process.pid)
         except Exception as e:
             return f"Error stopping training: {e}"
         finally:
             self.process.wait()
-        return "Training forcibly stopped"
+        return "Training forcibly stopped"
+    
+    def terminate_process_tree(pid):
+        try:
+            parent = psutil.Process(pid)
+            children = parent.children(recursive=True)  # Get child processes
+            for child in children:
+                child.terminate()
+            parent.terminate()
+        except psutil.NoSuchProcess:
+            pass

scripts/rename_keys.py

@@ -7,10 +7,10 @@
 def rename_keys(file, outfile: str)-> bool:
     sd, metadata = load_state_dict(file, torch.float32)
 
-    keys_to_normalize = [key for key in sd.keys()]
-    values_to_normalize = [sd[key].to(torch.float32) for key in keys_to_normalize]
+    keys_to_rename = [key for key in sd.keys()]
+    values = [sd[key].to(torch.float32) for key in keys_to_rename]
     new_sd = dict()
-    for key, value in zip(keys_to_normalize, values_to_normalize):
+    for key, value in zip(keys_to_rename, values):
         new_sd[key.replace("transformer.", "")] = value
 
     save_to_file(outfile, new_sd, torch.float16, metadata)

test/test_trainer_config_validator.py

@@ -2,6 +2,8 @@
 import pytest
 from unittest.mock import patch
 
+import yaml
+
 from trainer_config_validator import TrainerValidator
 
 @pytest.fixture
@@ -55,6 +57,23 @@ def test_valid_config(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_config_template():
+    config = None
+    with open('config/config_template.yaml', "r") as file:
+        config = yaml.safe_load(file)
+    config['path_to_finetrainers'] = '/path/to/finetrainers'
+    config['data_root'] = '/path/to/data'
+    config['pretrained_model_name_or_path'] = 'pretrained_model'
+    
+    trainer_validator = TrainerValidator(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_not_set(trainer_validator):
+    trainer_validator.config['data_root'] = ''
+    with pytest.raises(ValueError, match="data_root is required"):
+        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"):