From d8b3175b570a76bed6cc7018bbe90c3b629a09e0 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 13 Jun 2024 13:57:20 +0200 Subject: [PATCH 01/27] cleanup before sd3 integration --- main.py | 1 - trainer/models.py | 7 +++- trainer/optimizer.py | 40 ++++++++++++++++++- trainer/preprocess.py | 18 +++++---- trainer/utils/utils.py | 8 ++-- ..._face.json => training_args_face_sd15.json | 24 +++++------ training_args_face_sdxl.json | 27 +++++++++++++ training_args_object.json | 22 +++++----- training_args_style_sd15.json | 16 ++++---- training_args_style_sdxl.json | 20 +++++----- 10 files changed, 124 insertions(+), 59 deletions(-) rename training_args_face.json => training_args_face_sd15.json (51%) create mode 100644 training_args_face_sdxl.json diff --git a/main.py b/main.py index 6c38cc6..6c8be78 100755 --- a/main.py +++ b/main.py @@ -12,7 +12,6 @@ import torch import torch.utils.checkpoint from tqdm import tqdm -import prodigyopt from typing import Union, Iterable, List, Dict, Tuple, Optional, cast #from diffusers.training_utils import cast_training_params diff --git a/trainer/models.py b/trainer/models.py index f0d374a..3ccae20 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -6,8 +6,11 @@ from diffusers import AutoencoderKL, DDPMScheduler, EulerDiscreteScheduler, UNet ############################################################################################################ -SDXL_MODEL_CACHE = "./models/juggernaut_v6.safetensors" -SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernautXL_v6.safetensors" +#SDXL_MODEL_CACHE = "./models/juggernaut_v6.safetensors" +#SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernautXL_v6.safetensors" + +SDXL_MODEL_CACHE = "./models/sd_xl_base_1.0_0.9vae.safetensors" +SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/sd_xl_base_1.0_0.9vae.safetensors" #SDXL_MODEL_CACHE = "./models/Juggernaut-X-RunDiffusion-NSFW.safetensors" #SDXL_URL = "https://huggingface.co/RunDiffusion/Juggernaut-X-v10/resolve/main/Juggernaut-X-RunDiffusion-NSFW.safetensors" diff --git a/trainer/optimizer.py b/trainer/optimizer.py index f6ff6ff..e9a7b72 100644 --- a/trainer/optimizer.py +++ b/trainer/optimizer.py @@ -35,6 +35,39 @@ def get_unet_optimizer( print(f"Created {optimizer_name} optimizer for unet!") return optimizer_unet +# Taken (and slightly modified) from B-LoRA repo https://github.com/yardenfren1996/B-LoRA/blob/main/blora_utils.py +def is_belong_to_blocks(key, blocks): + try: + for g in blocks: + if g in key: + return True + return False + except Exception as e: + raise type(e)(f"failed to is_belong_to_block, due to: {e}") + +def get_unet_lora_target_modules(unet, use_blora, target_blocks=None): + if use_blora: + content_b_lora_blocks = "unet.up_blocks.0.attentions.0" + style_b_lora_blocks = "unet.up_blocks.0.attentions.1" + target_blocks = [content_b_lora_blocks, style_b_lora_blocks] + try: + blocks = [(".").join(blk.split(".")[1:]) for blk in target_blocks] + + attns = [ + attn_processor_name.rsplit(".", 1)[0] + for attn_processor_name, _ in unet.attn_processors.items() + if is_belong_to_blocks(attn_processor_name, blocks) + ] + + target_modules = [f"{attn}.{mat}" for mat in ["to_k", "to_q", "to_v", "to_out.0", "conv2"] for attn in attns] + return target_modules + except Exception as e: + raise type(e)( + f"failed to get_target_modules, due to: {e}. " + f"Please check the modules specified in --lora_unet_blocks are correct" + ) + + def get_unet_lora_parameters( lora_rank, lora_alpha_multiplier: float, @@ -43,12 +76,15 @@ def get_unet_lora_parameters( unet, pipe, ): + + #target_modules = get_unet_lora_target_modules(unet, use_blora=True) + target_modules = ["to_k", "to_q", "to_v", "to_out.0", "conv2"] + unet_lora_config = LoraConfig( r=lora_rank, lora_alpha=lora_rank * lora_alpha_multiplier, init_lora_weights="gaussian", - target_modules=["to_k", "to_q", "to_v", "to_out.0", "conv2"], - #target_modules=["conv1", "conv2", "norm1", "norm2", "proj_in"], # TODO grid-search params for sd15 + target_modules=target_modules, use_dora=use_dora, ) diff --git a/trainer/preprocess.py b/trainer/preprocess.py index 33dc996..d6c8697 100755 --- a/trainer/preprocess.py +++ b/trainer/preprocess.py @@ -777,7 +777,10 @@ def load_and_save_masks_and_captions( # Cleanup prompts using chatgpt: captions = [fix_prompt(caption) for caption in captions] - captions, trigger_text, gpt_concept_description = post_process_captions(captions, caption_text, concept_mode, seed) + trigger_text = "" + gpt_concept_description = None + if not config.remove_ti_token_from_prompts: + captions, trigger_text, gpt_concept_description = post_process_captions(captions, caption_text, concept_mode, seed) aug_imgs, aug_caps = [],[] # if we still have a very small amount of imgs, do some basic augmentation: @@ -854,12 +857,7 @@ def load_and_save_masks_and_captions( os.remove(os.path.join(output_dir, file)) os.makedirs(output_dir, exist_ok=True) - - # Make sure we've correctly inserted the TOK into every caption: - captions = ["TOK, " + caption if "TOK" not in caption else caption for caption in captions] - for caption in captions: - print(caption) - + if config.remove_ti_token_from_prompts: print('------------------ WARNING -------------------') print("Removing 'TOK, ' from captions...") @@ -871,6 +869,12 @@ def load_and_save_masks_and_captions( replace_str = "" captions = [caption.replace("TOK, ", replace_str + ", ") for caption in captions] captions = [caption.replace("TOK", replace_str) for caption in captions] + else: + captions = ["TOK, " + caption if "TOK" not in caption else caption for caption in captions] + + print("Final captions:") + for caption in captions: + print(caption) # iterate through the images, masks, and captions and add a row to the dataframe for each print("Saving final training dataset...") diff --git a/trainer/utils/utils.py b/trainer/utils/utils.py index f9337eb..369ead8 100755 --- a/trainer/utils/utils.py +++ b/trainer/utils/utils.py @@ -100,15 +100,17 @@ def print_system_info(): # Print disk space information disk_usage = psutil.disk_usage('/') - free_disk = disk_usage.free // (1024 * 1024) + total_disk = disk_usage.total // (1024 * 1024) + used_disk = disk_usage.used // (1024 * 1024) percent_disk_used = disk_usage.percent - print(f"Free disk space: {free_disk} MB with {percent_disk_used}% used") + print(f"Used disk space: {used_disk}/{total_disk} MB = {percent_disk_used}% used") # Print RAM information virtual_mem = psutil.virtual_memory() + total_ram = virtual_mem.total // (1024 * 1024) current_ram = virtual_mem.used // (1024 * 1024) percent_ram_used = virtual_mem.percent - print(f"Current used RAM: {current_ram} MB with {percent_ram_used}% used") + print(f"Current used RAM: {current_ram}/{total_ram} MB = {percent_ram_used}% used") except Exception as e: print(f'Error in gathering system info: {str(e)}') diff --git a/training_args_face.json b/training_args_face_sd15.json similarity index 51% rename from training_args_face.json rename to training_args_face_sd15.json index 8e653a2..6d07e0a 100644 --- a/training_args_face.json +++ b/training_args_face_sd15.json @@ -1,31 +1,27 @@ { - "output_dir": "lora_models/xander_sd15_final", + "output_dir": "lora_models/xander", "sd_model_version": "sd15", - "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_big.zip", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", "concept_mode": "face", "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 6, - "max_train_steps": 600, + "n_sample_imgs": 4, + "max_train_steps": 400, "token_warmup_steps": 0, - "checkpointing_steps": 100, - "gradient_accumulation_steps": 1, - - "sample_imgs_lora_scale": 0.8, - "n_tokens": 2, + "checkpointing_steps": 200, "ti_lr": 0.001, - "remove_ti_token_from_prompts": false, + "ti_weight_decay": 0.0005, + "remove_ti_token_from_prompts": false, "text_encoder_lora_optimizer": null, - "text_encoder_lora_lr": 0.5e-4, + "text_encoder_lora_lr": 1.0e-4, "text_encoder_lora_weight_decay": 1e-5, - "text_encoder_lora_rank": 16, + "text_encoder_lora_rank": 12, "unet_lr": 0.001, - "lora_alpha_multiplier": 1.0, "lora_rank": 16, "use_dora": false, - "caption_model": "gpt4-v", + "caption_model": "blip", "debug": true } \ No newline at end of file diff --git a/training_args_face_sdxl.json b/training_args_face_sdxl.json new file mode 100644 index 0000000..58fedb6 --- /dev/null +++ b/training_args_face_sdxl.json @@ -0,0 +1,27 @@ +{ + "output_dir": "lora_models/xander", + "sd_model_version": "sdxl", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", + "concept_mode": "face", + "seed": 0, + "resolution": 512, + "train_batch_size": 4, + "n_sample_imgs": 4, + "max_train_steps": 400, + "token_warmup_steps": 0, + "checkpointing_steps": 200, + "ti_lr": 0.001, + "ti_weight_decay": 0.0005, + + "remove_ti_token_from_prompts": false, + "text_encoder_lora_optimizer": null, + "text_encoder_lora_lr": 1.0e-4, + "text_encoder_lora_weight_decay": 1e-5, + "text_encoder_lora_rank": 12, + + "unet_lr": 0.001, + "lora_rank": 16, + "use_dora": false, + "caption_model": "blip", + "debug": true +} \ No newline at end of file diff --git a/training_args_object.json b/training_args_object.json index 3f1a3a1..6eef861 100644 --- a/training_args_object.json +++ b/training_args_object.json @@ -1,25 +1,27 @@ { - "output_dir": "lora_models/object", + "output_dir": "lora_models/banny", "sd_model_version": "sdxl", - "lora_training_urls": "/home/rednax/Documents/datasets/DOV/lizzo/full body", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/banny_best.zip", "concept_mode": "object", - "seed": 1, - "resolution": 640, + "seed": 0, + "resolution": 512, "train_batch_size": 4, "n_sample_imgs": 4, - "max_train_steps": 420, + "max_train_steps": 400, "token_warmup_steps": 0, - "checkpointing_steps": 60, - "gradient_accumulation_steps": 1, - "n_tokens": 2, + "checkpointing_steps": 200, "ti_lr": 0.001, "ti_weight_decay": 0.0005, + + "remove_ti_token_from_prompts": false, "text_encoder_lora_optimizer": null, "text_encoder_lora_lr": 1.0e-4, "text_encoder_lora_weight_decay": 1e-5, "text_encoder_lora_rank": 12, - "lora_rank": 12, + + "unet_lr": 0.001, + "lora_rank": 16, "use_dora": false, - "caption_model": "gpt4-v", + "caption_model": "blip", "debug": true } \ No newline at end of file diff --git a/training_args_style_sd15.json b/training_args_style_sd15.json index 1c12942..6e1d008 100644 --- a/training_args_style_sd15.json +++ b/training_args_style_sd15.json @@ -1,17 +1,15 @@ { - "output_dir": "lora_models/does_best", + "output_dir": "lora_models/clipx", "sd_model_version": "sd15", - "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/does.zip", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/clipx_tiny.zip", "concept_mode": "style", - "seed": 1, - "resolution": 640, + "seed": 0, + "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 6, - "max_train_steps": 600, + "n_sample_imgs": 4, + "max_train_steps": 400, "token_warmup_steps": 0, - "checkpointing_steps": 100, - "gradient_accumulation_steps": 1, - "n_tokens": 2, + "checkpointing_steps": 200, "ti_lr": 0.001, "ti_weight_decay": 0.0005, diff --git a/training_args_style_sdxl.json b/training_args_style_sdxl.json index 30cc5f1..14150be 100644 --- a/training_args_style_sdxl.json +++ b/training_args_style_sdxl.json @@ -1,29 +1,27 @@ { - "output_dir": "lora_models/Journey", + "output_dir": "lora_models/clipx", "sd_model_version": "sdxl", - "lora_training_urls": "/home/rednax/Documents/datasets/journey", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/clipx_tiny.zip", "concept_mode": "style", "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 6, - "max_train_steps": 1000, + "n_sample_imgs": 4, + "max_train_steps": 400, "token_warmup_steps": 0, - "checkpointing_steps": 100, - "gradient_accumulation_steps": 1, - "n_tokens": 2, + "checkpointing_steps": 200, "ti_lr": 0.001, "ti_weight_decay": 0.0005, + + "remove_ti_token_from_prompts": false, "text_encoder_lora_optimizer": null, "text_encoder_lora_lr": 1.0e-4, "text_encoder_lora_weight_decay": 1e-5, "text_encoder_lora_rank": 12, "unet_lr": 0.001, - "prodigy_d_coef": 1.0, - "unet_prodigy_growth_factor": 1.05, "lora_rank": 16, - "use_dora": true, - "caption_model": "gpt4-v", + "use_dora": false, + "caption_model": "blip", "debug": true } \ No newline at end of file From b02221b23c4b4c15c1200086aee2b9869e459f4e Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 13 Jun 2024 14:24:32 +0200 Subject: [PATCH 02/27] cleanup before sd3 integration --- main.py | 2 +- test_trainer.sh | 8 ++++++++ .../training_args_face_sd15.json | 8 ++++---- .../training_args_face_sdxl.json | 6 +++--- .../training_args_object.json | 4 ++-- .../training_args_style_sd15.json | 4 ++-- .../training_args_style_sdxl.json | 4 ++-- trainer/models.py | 14 ++++---------- 8 files changed, 26 insertions(+), 24 deletions(-) create mode 100644 test_trainer.sh rename training_args_face_sd15.json => train_configs/training_args_face_sd15.json (83%) rename training_args_face_sdxl.json => train_configs/training_args_face_sdxl.json (89%) rename training_args_object.json => train_configs/training_args_object.json (91%) rename training_args_style_sd15.json => train_configs/training_args_style_sd15.json (91%) rename training_args_style_sdxl.json => train_configs/training_args_style_sdxl.json (91%) diff --git a/main.py b/main.py index 6c8be78..326538d 100755 --- a/main.py +++ b/main.py @@ -368,7 +368,7 @@ def train( token_stds[f'text_encoder_{idx}'][std_i].append(embedding_stds[std_i].item()) # Print some statistics: - if config.debug and (global_step % config.checkpointing_steps == 0) and (global_step < (config.max_train_steps - 25)) and global_step > -1: + if config.debug and (global_step % config.checkpointing_steps == 0) and (global_step < (config.max_train_steps - 25)) and global_step > 0: output_save_dir = f"{checkpoint_dir}/checkpoint-{global_step}" os.makedirs(output_save_dir, exist_ok=True) diff --git a/test_trainer.sh b/test_trainer.sh new file mode 100644 index 0000000..831bcb9 --- /dev/null +++ b/test_trainer.sh @@ -0,0 +1,8 @@ +# Set GPU ID to run these jobs on: +GPU_ID="device=0" + +python main.py train_configs/training_args_face_sdxl.json +python main.py train_configs/training_args_face_sd15.json +python main.py train_configs/training_args_object.json +python main.py train_configs/training_args_style_sd15.json +python main.py train_configs/training_args_style_sdxl.json diff --git a/training_args_face_sd15.json b/train_configs/training_args_face_sd15.json similarity index 83% rename from training_args_face_sd15.json rename to train_configs/training_args_face_sd15.json index 6d07e0a..dc7f423 100644 --- a/training_args_face_sd15.json +++ b/train_configs/training_args_face_sd15.json @@ -1,15 +1,15 @@ { - "output_dir": "lora_models/xander", + "output_dir": "lora_models/test_xander", "sd_model_version": "sd15", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", "concept_mode": "face", "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 4, - "max_train_steps": 400, + "n_sample_imgs": 6, + "max_train_steps": 600, "token_warmup_steps": 0, - "checkpointing_steps": 200, + "checkpointing_steps": 300, "ti_lr": 0.001, "ti_weight_decay": 0.0005, diff --git a/training_args_face_sdxl.json b/train_configs/training_args_face_sdxl.json similarity index 89% rename from training_args_face_sdxl.json rename to train_configs/training_args_face_sdxl.json index 58fedb6..3662f3d 100644 --- a/training_args_face_sdxl.json +++ b/train_configs/training_args_face_sdxl.json @@ -1,12 +1,12 @@ { - "output_dir": "lora_models/xander", + "output_dir": "lora_models/test_xander", "sd_model_version": "sdxl", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", "concept_mode": "face", - "seed": 0, + "seed": 1, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 4, + "n_sample_imgs": 6, "max_train_steps": 400, "token_warmup_steps": 0, "checkpointing_steps": 200, diff --git a/training_args_object.json b/train_configs/training_args_object.json similarity index 91% rename from training_args_object.json rename to train_configs/training_args_object.json index 6eef861..a018a9b 100644 --- a/training_args_object.json +++ b/train_configs/training_args_object.json @@ -1,12 +1,12 @@ { - "output_dir": "lora_models/banny", + "output_dir": "lora_models/test_banny", "sd_model_version": "sdxl", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/banny_best.zip", "concept_mode": "object", "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 4, + "n_sample_imgs": 6, "max_train_steps": 400, "token_warmup_steps": 0, "checkpointing_steps": 200, diff --git a/training_args_style_sd15.json b/train_configs/training_args_style_sd15.json similarity index 91% rename from training_args_style_sd15.json rename to train_configs/training_args_style_sd15.json index 6e1d008..10ad4d8 100644 --- a/training_args_style_sd15.json +++ b/train_configs/training_args_style_sd15.json @@ -1,12 +1,12 @@ { - "output_dir": "lora_models/clipx", + "output_dir": "lora_models/test_clipx", "sd_model_version": "sd15", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/clipx_tiny.zip", "concept_mode": "style", "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 4, + "n_sample_imgs": 6, "max_train_steps": 400, "token_warmup_steps": 0, "checkpointing_steps": 200, diff --git a/training_args_style_sdxl.json b/train_configs/training_args_style_sdxl.json similarity index 91% rename from training_args_style_sdxl.json rename to train_configs/training_args_style_sdxl.json index 14150be..e4ca178 100644 --- a/training_args_style_sdxl.json +++ b/train_configs/training_args_style_sdxl.json @@ -1,12 +1,12 @@ { - "output_dir": "lora_models/clipx", + "output_dir": "lora_models/test_clipx", "sd_model_version": "sdxl", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/clipx_tiny.zip", "concept_mode": "style", "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 4, + "n_sample_imgs": 6, "max_train_steps": 400, "token_warmup_steps": 0, "checkpointing_steps": 200, diff --git a/trainer/models.py b/trainer/models.py index 3ccae20..959961d 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -6,21 +6,15 @@ from diffusers import AutoencoderKL, DDPMScheduler, EulerDiscreteScheduler, UNet ############################################################################################################ -#SDXL_MODEL_CACHE = "./models/juggernaut_v6.safetensors" -#SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernautXL_v6.safetensors" +SDXL_MODEL_CACHE = "./models/juggernaut_v6.safetensors" +SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernautXL_v6.safetensors" -SDXL_MODEL_CACHE = "./models/sd_xl_base_1.0_0.9vae.safetensors" -SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/sd_xl_base_1.0_0.9vae.safetensors" - -#SDXL_MODEL_CACHE = "./models/Juggernaut-X-RunDiffusion-NSFW.safetensors" -#SDXL_URL = "https://huggingface.co/RunDiffusion/Juggernaut-X-v10/resolve/main/Juggernaut-X-RunDiffusion-NSFW.safetensors" +#SDXL_MODEL_CACHE = "./models/sd_xl_base_1.0_0.9vae.safetensors" +#SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/comfyui/models2/checkpoints/sd_xl_base_1.0_0.9vae.safetensors" SD15_MODEL_CACHE = "./models/juggernaut_reborn.safetensors" SD15_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernaut_reborn.safetensors" -#SD15_MODEL_CACHE = "./models/DreamShaper_6.31_BakedVae.safetensors" -#SD15_URL = "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaper_6.31_BakedVae.safetensors" - #SD15_MODEL_CACHE = "./models/photon_v1.safetensors" #SD15_URL = "https://civitai.com/api/download/models/90072" From faf864ce15794f6e1cd254cccbd84ed3a111ad7e Mon Sep 17 00:00:00 2001 From: xander Date: Fri, 5 Jul 2024 17:07:33 +0200 Subject: [PATCH 03/27] tiny tweak to learning rates for SDXL --- main.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/main.py b/main.py index 326538d..3a74a5b 100755 --- a/main.py +++ b/main.py @@ -216,9 +216,9 @@ def train( # default value of cold (pre-warmup) optimizer lr: if config.sd_model_version == "sdxl": # let textual_inversion do the work first! - base_lr = 0.5e-5 + base_lr = 1.0e-5 elif config.sd_model_version == "sd15": - # let lora training kick in soonish + # let lora training kick in soonish (pure ti for sd15 is not working super well in my tests) base_lr = 1.0e-4 ####################################################################################################### From ef6635fafbd1ceb012552627efa2ed14b42f6aa0 Mon Sep 17 00:00:00 2001 From: xander Date: Fri, 5 Jul 2024 17:11:06 +0200 Subject: [PATCH 04/27] update train cmd --- README.md | 2 +- train_configs/training_args_object.json | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index a10b8fa..12dacdd 100755 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ Install all dependencies using then you can simply run: -`python main.py -c training_args.json` +`python main.py train_configs/training_args.json` to start a training job. Adjust the arguments inside `training_args.json` to setup a custom training job. diff --git a/train_configs/training_args_object.json b/train_configs/training_args_object.json index a018a9b..63d49f4 100644 --- a/train_configs/training_args_object.json +++ b/train_configs/training_args_object.json @@ -1,13 +1,13 @@ { "output_dir": "lora_models/test_banny", - "sd_model_version": "sdxl", + "sd_model_version": "sd15", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/banny_best.zip", "concept_mode": "object", "seed": 0, "resolution": 512, "train_batch_size": 4, "n_sample_imgs": 6, - "max_train_steps": 400, + "max_train_steps": 800, "token_warmup_steps": 0, "checkpointing_steps": 200, "ti_lr": 0.001, From fdd531d3c5ee222f675d473363fa688525b8117a Mon Sep 17 00:00:00 2001 From: xander Date: Sat, 6 Jul 2024 17:42:49 +0200 Subject: [PATCH 05/27] fix full finetuning and add 8bitadam --- main.py | 3 ++- requirements.txt | 3 ++- train_configs/training_args_beeple.json | 29 +++++++++++++++++++++++++ train_configs/training_args_object.json | 7 +++--- trainer/config.py | 2 +- trainer/optimizer.py | 4 +++- 6 files changed, 41 insertions(+), 7 deletions(-) create mode 100644 train_configs/training_args_beeple.json diff --git a/main.py b/main.py index 3a74a5b..546e220 100755 --- a/main.py +++ b/main.py @@ -123,6 +123,7 @@ def train( print(f"Doing full fine-tuning on the U-Net") unet.requires_grad_(True) unet_lora_parameters = None + optimizer_text_encoder_lora = None unet_trainable_params = unet.parameters() else: # Do lora-training instead. @@ -319,7 +320,7 @@ def train( loss += 0.0 * concept_description_loss losses['concept_description_loss'].append(concept_description_loss.item()) - if config.l1_penalty > 0.0: + if config.l1_penalty > 0.0 and unet_lora_parameters: # Compute normalized L1 norm (mean of abs sum) of all lora parameters: l1_norm = sum(p.abs().sum() for p in unet_lora_parameters) / sum(p.numel() for p in unet_lora_parameters) loss += config.l1_penalty * l1_norm diff --git a/requirements.txt b/requirements.txt index b1ecfdc..4faf486 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,4 +14,5 @@ openai>=1.14.0 python-dotenv prodigyopt omegaconf -ujson \ No newline at end of file +ujson +bitsandbytes \ No newline at end of file diff --git a/train_configs/training_args_beeple.json b/train_configs/training_args_beeple.json new file mode 100644 index 0000000..cc7bf63 --- /dev/null +++ b/train_configs/training_args_beeple.json @@ -0,0 +1,29 @@ +{ + "output_dir": "lora_models/beeple", + "sd_model_version": "sdxl", + "lora_training_urls": "/home/rednax/SSD2TB/Github_repos/Eden/images/beeple_small", + "concept_mode": "style", + "seed": 0, + "resolution": 512, + "train_batch_size": 2, + "n_sample_imgs": 4, + "max_train_steps": 400, + "token_warmup_steps": 0, + "checkpointing_steps": 100, + "ti_lr": 0.001, + "ti_weight_decay": 0.0005, + + "remove_ti_token_from_prompts": false, + "text_encoder_lora_optimizer": null, + "text_encoder_lora_lr": 1.0e-4, + "text_encoder_lora_weight_decay": 1e-5, + "text_encoder_lora_rank": 12, + + "unet_lr": 0.001, + "lora_rank": 4, + "use_dora": false, + "is_lora": false, + "unet_optimizer_type": "AdamW8bit", + "caption_model": "blip", + "debug": true +} \ No newline at end of file diff --git a/train_configs/training_args_object.json b/train_configs/training_args_object.json index 63d49f4..4334731 100644 --- a/train_configs/training_args_object.json +++ b/train_configs/training_args_object.json @@ -1,8 +1,9 @@ { - "output_dir": "lora_models/test_banny", + "output_dir": "lora_models/test_xander", "sd_model_version": "sd15", - "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/banny_best.zip", - "concept_mode": "object", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander.zip", + "concept_mode": "face", + "sample_imgs_lora_scale": 0.8, "seed": 0, "resolution": 512, "train_batch_size": 4, diff --git a/trainer/config.py b/trainer/config.py index 63ac61a..8dbcd22 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -25,7 +25,7 @@ class TrainingConfig(BaseModel): gradient_accumulation_steps: int = 1 is_lora: bool = True - unet_optimizer_type: Literal["adamw", "prodigy"] = "adamw" + unet_optimizer_type: Literal["adamw", "prodigy", "AdamW8bit"] = "adamw" unet_lr_warmup_steps: int = None # slowly increase the learning rate of the adamw unet optimizer unet_lr: float = 1.0e-3 prodigy_d_coef: float = 1.0 diff --git a/trainer/optimizer.py b/trainer/optimizer.py index e9a7b72..ba1f8af 100644 --- a/trainer/optimizer.py +++ b/trainer/optimizer.py @@ -15,7 +15,9 @@ def get_unet_optimizer( if optimizer_name == "adamw": optimizer_unet = torch.optim.AdamW(unet_trainable_params, lr = 1e-4, weight_decay=lora_weight_decay if not use_dora else 0.0) - + elif optimizer_name == "AdamW8bit": + import bitsandbytes as bnb + optimizer_unet = bnb.optim.AdamW8bit(unet_trainable_params, lr = 1e-4, weight_decay=lora_weight_decay) elif optimizer_name == "prodigy": # Note: the specific settings of Prodigy seem to matter A LOT optimizer_unet = prodigyopt.Prodigy( From 892cf61c528de58c38440fed82899a4a715cc495 Mon Sep 17 00:00:00 2001 From: xander Date: Sat, 6 Jul 2024 18:43:03 +0200 Subject: [PATCH 06/27] add disable_ti flag --- main.py | 23 ++++++++++------------- train_configs/training_args_beeple.json | 14 +++++++------- trainer/config.py | 2 +- trainer/preprocess.py | 6 +++--- 4 files changed, 21 insertions(+), 24 deletions(-) diff --git a/main.py b/main.py index 546e220..b097d33 100755 --- a/main.py +++ b/main.py @@ -112,12 +112,16 @@ def train( embedding_handler.make_embeddings_trainable() - optimizer_ti, textual_inversion_params = get_textual_inversion_optimizer( - text_encoders=text_encoders, - textual_inversion_lr=config.ti_lr, - textual_inversion_weight_decay=config.ti_weight_decay, - optimizer_name=config.ti_optimizer ## hardcoded - ) + if not config.disable_ti: + optimizer_ti, textual_inversion_params = get_textual_inversion_optimizer( + text_encoders=text_encoders, + textual_inversion_lr=config.ti_lr, + textual_inversion_weight_decay=config.ti_weight_decay, + optimizer_name=config.ti_optimizer ## hardcoded + ) + else: + optimizer_ti = None + textual_inversion_params = None if not config.is_lora: # This code pathway has not been tested in a long while print(f"Doing full fine-tuning on the U-Net") @@ -128,7 +132,6 @@ def train( else: # Do lora-training instead. # https://huggingface.co/docs/peft/main/en/developer_guides/lora#rank-stabilized-lora - # target_blocks=["block"] for original IP-Adapter # target_blocks=["up_blocks.0.attentions.1"] for style blocks only # target_blocks = ["up_blocks.0.attentions.1", "down_blocks.2.attentions.1"] # for style+layout blocks @@ -349,12 +352,6 @@ def train( grad_norms[f'text_encoder_{i}'].append(text_encoder_norm) optimizer_collection.step() - - # after every optimizer step, we do some manual intervention of the embeddings to regularize them: - if optimizer_collection.get_lr('textual_inversion') > 0.0: - #embedding_handler.fix_embedding_std(config.off_ratio_power) - pass - optimizer_collection.zero_grad() ############################################################################################################# diff --git a/train_configs/training_args_beeple.json b/train_configs/training_args_beeple.json index cc7bf63..8a57378 100644 --- a/train_configs/training_args_beeple.json +++ b/train_configs/training_args_beeple.json @@ -5,21 +5,21 @@ "concept_mode": "style", "seed": 0, "resolution": 512, - "train_batch_size": 2, - "n_sample_imgs": 4, - "max_train_steps": 400, + "train_batch_size": 4, + "n_sample_imgs": 6, + "max_train_steps": 1000, "token_warmup_steps": 0, - "checkpointing_steps": 100, - "ti_lr": 0.001, + "checkpointing_steps": 200, + "ti_lr": 0.000, "ti_weight_decay": 0.0005, - "remove_ti_token_from_prompts": false, + "disable_ti": true, "text_encoder_lora_optimizer": null, "text_encoder_lora_lr": 1.0e-4, "text_encoder_lora_weight_decay": 1e-5, "text_encoder_lora_rank": 12, - "unet_lr": 0.001, + "unet_lr": 0.0003, "lora_rank": 4, "use_dora": false, "is_lora": false, diff --git a/trainer/config.py b/trainer/config.py index 8dbcd22..6e01426 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -62,7 +62,7 @@ class TrainingConfig(BaseModel): output_dir: str = "lora_models/unnamed" debug: bool = False allow_tf32: bool = True - remove_ti_token_from_prompts: bool = False + disable_ti: bool = False weight_type: Literal["fp16", "bf16", "fp32"] = "bf16" n_tokens: int = 2 inserting_list_tokens: List[str] = ["",""] diff --git a/trainer/preprocess.py b/trainer/preprocess.py index d6c8697..74d61f5 100755 --- a/trainer/preprocess.py +++ b/trainer/preprocess.py @@ -779,7 +779,7 @@ def load_and_save_masks_and_captions( captions = [fix_prompt(caption) for caption in captions] trigger_text = "" gpt_concept_description = None - if not config.remove_ti_token_from_prompts: + if not config.disable_ti: captions, trigger_text, gpt_concept_description = post_process_captions(captions, caption_text, concept_mode, seed) aug_imgs, aug_caps = [],[] @@ -858,10 +858,10 @@ def load_and_save_masks_and_captions( os.makedirs(output_dir, exist_ok=True) - if config.remove_ti_token_from_prompts: + if config.disable_ti: print('------------------ WARNING -------------------') print("Removing 'TOK, ' from captions...") - print("This will completely break textual_inversion!!") + print("This will completely disable textual_inversion!!") print('------------------ WARNING -------------------') if gpt_concept_description: replace_str = gpt_concept_description From 4fde1b7dd62190802cafc6bad00163513861670c Mon Sep 17 00:00:00 2001 From: xander Date: Sat, 6 Jul 2024 18:45:38 +0200 Subject: [PATCH 07/27] small tweaks --- trainer/checkpoint.py | 1 + 1 file changed, 1 insertion(+) diff --git a/trainer/checkpoint.py b/trainer/checkpoint.py index 485ba8e..16156df 100644 --- a/trainer/checkpoint.py +++ b/trainer/checkpoint.py @@ -187,6 +187,7 @@ def save_checkpoint( output_filename=os.path.join(output_dir, f"{name}.safetensors") ) else: + # Save the entire, finetuned unet weights: unet.save_pretrained(save_directory = output_dir) def load_checkpoint( From e1813c75d1de82fcfeb51fd4f7b7b1f3e8613280 Mon Sep 17 00:00:00 2001 From: xander Date: Sat, 6 Jul 2024 22:18:47 +0200 Subject: [PATCH 08/27] updates to avoid OOM when plotting hist --- main.py | 12 ++++++--- scripts/create_hyperparam_sweep.py | 35 ++++++++++++++----------- train_configs/training_args_beeple.json | 11 ++++---- trainer/checkpoint.py | 2 +- trainer/optimizer.py | 1 + trainer/utils/utils.py | 7 +++++ 6 files changed, 43 insertions(+), 25 deletions(-) diff --git a/main.py b/main.py index b097d33..f223385 100755 --- a/main.py +++ b/main.py @@ -219,8 +219,11 @@ def train( # default value of cold (pre-warmup) optimizer lr: if config.sd_model_version == "sdxl": - # let textual_inversion do the work first! - base_lr = 1.0e-5 + if config.is_lora: + # let textual_inversion do the work first! + base_lr = 1.0e-5 + else: + base_lr = 1.0e-4 elif config.sd_model_version == "sd15": # let lora training kick in soonish (pure ti for sd15 is not working super well in my tests) base_lr = 1.0e-4 @@ -400,7 +403,9 @@ def train( plot_torch_hist(embedding, global_step, os.path.join(config.output_dir, 'ti_embeddings') , f"enc_{idx}_tokid_{i}: {token}", min_val=-0.05, max_val=0.05, ymax_f = 0.05, color = 'red') embedding_handler.print_token_info() - plot_torch_hist(unet_lora_parameters if config.is_lora else unet.parameters(), global_step, config.output_dir, "lora_weights", min_val=-0.4, max_val=0.4, ymax_f = 0.08) + print("plotting...") + if config.is_lora: # plotting this hist for full unet parameters can run OOM + plot_torch_hist(unet_lora_parameters, global_step, config.output_dir, "lora_weights", min_val=-0.4, max_val=0.4, ymax_f = 0.08) plot_loss(losses, save_path=f'{config.output_dir}/losses.png') target_std_dict = {f"text_encoder_{idx}_target": embedding_handler.embeddings_settings[f"std_token_embedding_{idx}"].item() for idx in range(len(text_encoders)) if text_encoders[idx] is not None} plot_token_stds(token_stds, save_path=f'{config.output_dir}/token_stds.png', target_value_dict=target_std_dict) @@ -408,6 +413,7 @@ def train( plot_lrs(optimizer_collection.learning_rate_tracker, save_path=f'{config.output_dir}/learning_rates.png') plot_curve(prompt_embeds_norms, 'steps', 'norm', 'prompt_embed norms', save_path=f'{config.output_dir}/prompt_embeds_norms.png') + print("Getting ready to render validation images...") validation_prompts = render_images( pipe = pipe, render_size = config.validation_img_size, diff --git a/scripts/create_hyperparam_sweep.py b/scripts/create_hyperparam_sweep.py index af46f65..552fe74 100644 --- a/scripts/create_hyperparam_sweep.py +++ b/scripts/create_hyperparam_sweep.py @@ -48,43 +48,46 @@ output_sh_path = f"gridsearch_configs/{exp_name}.sh" hyperparameters = { "output_dir": [f"lora_models/{exp_name}"], - "sd_model_version": ["sd15", "sdxl"], + "sd_model_version": ["sdxl"], "lora_training_urls": [ - "/home/rednax/Documents/datasets/grimes" + "/home/rednax/SSD2TB/Github_repos/Eden/images/beeple_large", + "/home/rednax/SSD2TB/Github_repos/Eden/images/beeple" ], - "concept_mode": ['face'], + "concept_mode": ['style'], + "sample_imgs_lora_scale": [0.9], + "disable_ti": ['false', 'true'], "seed": [0], "resolution": [512], "train_batch_size": [4], - "n_sample_imgs": [6], - "max_train_steps": [400,800], - "checkpointing_steps": [100], + "n_sample_imgs": [8], + "max_train_steps": [1200], + "checkpointing_steps": [200], "gradient_accumulation_steps": [1], "n_tokens": [2], - "ti_lr": [0.001,0.0005], - "ti_weight_decay": [0.001,0.0], + "ti_lr": [0.001], + "ti_weight_decay": [0.001], "l1_penalty": [0.0], - "token_warmup_steps": [0,60], + "token_warmup_steps": [0], "tok_cov_reg_w": [2000], - "cond_reg_w": [0.01e-5], - "tok_cond_reg_w": [0.01e-5], - "unet_prodigy_growth_factor": [1.05], - "unet_lr": [0.001], + "unet_lr": [0.001, 0.0001], "lora_alpha_multiplier": [1.0], "prodigy_d_coef": [1.0], "lora_weight_decay": [0.001], - "lora_rank": [16,32], - "use_dora": ['false', 'true'], + "lora_rank": [16], + "use_dora": ['false'], + + "unet_optimizer_type": ['AdamW8bit'], + "is_lora": ['false'], "text_encoder_lora_optimizer": [None], "text_encoder_lora_lr": [0.0e-4], "snr_gamma": [5.0], "caption_model": ["blip", "gpt4-v"], - "augment_imgs_up_to_n": [20,40], + "augment_imgs_up_to_n": [40], "verbose": ['true'], "debug": ['true'] } diff --git a/train_configs/training_args_beeple.json b/train_configs/training_args_beeple.json index 8a57378..e53dcad 100644 --- a/train_configs/training_args_beeple.json +++ b/train_configs/training_args_beeple.json @@ -1,16 +1,17 @@ { "output_dir": "lora_models/beeple", "sd_model_version": "sdxl", - "lora_training_urls": "/home/rednax/SSD2TB/Github_repos/Eden/images/beeple_small", + "lora_training_urls": "/home/rednax/SSD2TB/Github_repos/Eden/images/beeple_large", "concept_mode": "style", + "sample_imgs_lora_scale": 0.9, "seed": 0, "resolution": 512, "train_batch_size": 4, - "n_sample_imgs": 6, - "max_train_steps": 1000, + "n_sample_imgs": 8, + "max_train_steps": 1200, "token_warmup_steps": 0, "checkpointing_steps": 200, - "ti_lr": 0.000, + "ti_lr": 0.001, "ti_weight_decay": 0.0005, "disable_ti": true, @@ -19,7 +20,7 @@ "text_encoder_lora_weight_decay": 1e-5, "text_encoder_lora_rank": 12, - "unet_lr": 0.0003, + "unet_lr": 0.0001, "lora_rank": 4, "use_dora": false, "is_lora": false, diff --git a/trainer/checkpoint.py b/trainer/checkpoint.py index 16156df..1789287 100644 --- a/trainer/checkpoint.py +++ b/trainer/checkpoint.py @@ -145,7 +145,7 @@ def save_checkpoint( output_dir, "special_params.json" ) ) - + if is_lora: assert len(unet_lora_parameters) > 0, f"Expected len(unet_lora_parameters) to be greater than zero if is_lora is True" diff --git a/trainer/optimizer.py b/trainer/optimizer.py index ba1f8af..c3ff5f6 100644 --- a/trainer/optimizer.py +++ b/trainer/optimizer.py @@ -13,6 +13,7 @@ def get_unet_optimizer( ): ## unet_trainable_params can be unet.parameters() or a list of lora params + # These learning rates will get overwritten in main.py: if optimizer_name == "adamw": optimizer_unet = torch.optim.AdamW(unet_trainable_params, lr = 1e-4, weight_decay=lora_weight_decay if not use_dora else 0.0) elif optimizer_name == "AdamW8bit": diff --git a/trainer/utils/utils.py b/trainer/utils/utils.py index 369ead8..87a3886 100755 --- a/trainer/utils/utils.py +++ b/trainer/utils/utils.py @@ -124,6 +124,13 @@ def plot_torch_hist(parameters, step, checkpoint_dir, name, bins=100, min_val=-1 # Flatten and concatenate all parameters into a single tensor all_params = torch.cat([p.data.view(-1) for p in parameters]) + + # count number of parameters: + n_params = len(all_params) + + if n_params == 0 or n_params > 1e9: + return + norm = torch.norm(all_params) # Convert to CPU for plotting From b5cf857b1f21e4fbf634b7793db13ff96b79044f Mon Sep 17 00:00:00 2001 From: xander Date: Tue, 9 Jul 2024 18:33:52 +0200 Subject: [PATCH 09/27] push small changes --- main.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/main.py b/main.py index f223385..13c51b8 100755 --- a/main.py +++ b/main.py @@ -219,11 +219,10 @@ def train( # default value of cold (pre-warmup) optimizer lr: if config.sd_model_version == "sdxl": - if config.is_lora: - # let textual_inversion do the work first! + if config.is_lora: # let textual_inversion do the work first! base_lr = 1.0e-5 else: - base_lr = 1.0e-4 + base_lr = 3.0e-5 elif config.sd_model_version == "sd15": # let lora training kick in soonish (pure ti for sd15 is not working super well in my tests) base_lr = 1.0e-4 From 5abed1487fac85d2398f2238cc94146c2d7953a3 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 11 Jul 2024 13:36:11 +0200 Subject: [PATCH 10/27] setup lora_scale automation for validation grid --- trainer/config.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/trainer/config.py b/trainer/config.py index 6e01426..1005fd8 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -74,7 +74,7 @@ class TrainingConfig(BaseModel): unet_learning_rate: float = 1.0 lr_num_cycles: int = 1 lr_power: float = 1.0 - sample_imgs_lora_scale: float = 0.65 # Default lora scale for sampling the validation images + sample_imgs_lora_scale: float = None # Default lora scale for sampling the validation images dataloader_num_workers: int = 0 training_attributes: dict = {} aspect_ratio_bucketing: bool = False @@ -116,6 +116,12 @@ class TrainingConfig(BaseModel): self.left_right_flip_augmentation = False # always disable lr flips for face mode! self.mask_target_prompts = "face" #self.use_face_detection_instead = True + + if not self.sample_imgs_lora_scale: + if self.sd_model_version == "sdxl": + self.sample_imgs_lora_scale = 0.7 + else: + self.sample_imgs_lora_scale = 0.85 if self.use_dora: print(f"Disabling L1 penalty and LoRA weight decay for DORA training.") From 340336b53d26ab69b210f2f8f6b3992395a2a3bb Mon Sep 17 00:00:00 2001 From: aiXander Date: Thu, 11 Jul 2024 06:11:33 -0700 Subject: [PATCH 11/27] print config pre training start --- main.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/main.py b/main.py index 13c51b8..36406a6 100755 --- a/main.py +++ b/main.py @@ -543,6 +543,10 @@ if __name__ == "__main__": args = parser.parse_args() config = TrainingConfig.from_json(file_path=args.config_filename) + + print("Starting new LoRa training run with config:") + print(config) + for progress in train(config=config): print(f"Progress: {(100*progress):.2f}%", end="\r") From 7d8c845765a6133e6a47a300d149d3ee1ce349b0 Mon Sep 17 00:00:00 2001 From: aiXander Date: Fri, 12 Jul 2024 05:16:06 -0700 Subject: [PATCH 12/27] update yaml and print deps in config --- cog.yaml | 16 +++++++++------- trainer/config.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 7 deletions(-) diff --git a/cog.yaml b/cog.yaml index e284eba..4ee180c 100755 --- a/cog.yaml +++ b/cog.yaml @@ -3,16 +3,18 @@ build: gpu: true - cuda: "11.8" - python_version: "3.9" + cuda: "12.1" + python_version: "3.11" system_packages: - - "ffmpeg" - - "libgl1-mesa-glx" - - "libegl1-mesa-dev" - - "libsm6" - - "libxext6" + - "libcudnn8=8.9.0.131-1+cuda12.1" + - "libcudnn8-dev=8.9.0.131-1+cuda12.1" + python_requirements: requirements.txt run: + - apt-mark showhold + - apt-get update -y -qq + - apt-get upgrade -y --allow-change-held-packages + - apt-get install -y --allow-change-held-packages ffmpeg libgl1-mesa-glx libegl1-mesa-dev libsm6 libxext6 libglib2.0-0 - wget http://thegiflibrary.tumblr.com/post/11565547760 -O face_landmarker_v2_with_blendshapes.task -q https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task predict: "predict.py:Predictor" diff --git a/trainer/config.py b/trainer/config.py index 1005fd8..c8b3e8f 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -5,6 +5,53 @@ import json, time, os from typing import Literal from trainer.models import pretrained_models from trainer.utils.utils import pick_best_gpu_id +import subprocess + +def run_command(command): + try: + result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True) + return result.stdout + except subprocess.CalledProcessError as e: + return f"Error: {e.stderr}" + + +print("=================================================================") +print("=================================================================") + +print("NVIDIA Driver Version:") +nvidia_smi_output = run_command("nvidia-smi --query-gpu=driver_version --format=csv,noheader") +print(nvidia_smi_output.strip()) # strip() removes any trailing newlines + +print("\nNVCC Version:") +print(run_command("nvcc --version")) + +print("\nCUDNN Libraries:") +print(run_command("ldconfig -p | grep cudnn")) + +print("\nFull NVIDIA-SMI Output:") +print(run_command("nvidia-smi")) + +# Run nvcc --version +nvcc_version = run_command("nvcc --version") +print("NVCC Version:") +print(nvcc_version) + +# Run ldconfig -p | grep cudnn +cudnn_info = run_command("ldconfig -p | grep cudnn") +print("\nCUDNN Libraries:") +print(cudnn_info) + +# Check if CUDA is available +import torch +if torch.cuda.is_available(): + print(f"Torch version: {torch.__version__}") + print(f"CUDA Version: {torch.version.cuda}") + print(f"cuDNN Version: {torch.backends.cudnn.version()}") +else: + print("CUDA is not available.") + +print("-----------------------------------------------------------------") +print("-----------------------------------------------------------------") class TrainingConfig(BaseModel): lora_training_urls: str From 2f5eaeba7a0c1bfdec3d9a757527911f1eca6a92 Mon Sep 17 00:00:00 2001 From: Gene Kogan Date: Tue, 16 Jul 2024 03:44:35 -0700 Subject: [PATCH 13/27] Update cog.yaml --- cog.yaml | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/cog.yaml b/cog.yaml index 4ee180c..bf8d7b3 100755 --- a/cog.yaml +++ b/cog.yaml @@ -6,16 +6,11 @@ build: cuda: "12.1" python_version: "3.11" system_packages: - - "libcudnn8=8.9.0.131-1+cuda12.1" - - "libcudnn8-dev=8.9.0.131-1+cuda12.1" + - "ffmpeg" python_requirements: requirements.txt run: - - apt-mark showhold - - apt-get update -y -qq - - apt-get upgrade -y --allow-change-held-packages - - apt-get install -y --allow-change-held-packages ffmpeg libgl1-mesa-glx libegl1-mesa-dev libsm6 libxext6 libglib2.0-0 - - wget http://thegiflibrary.tumblr.com/post/11565547760 -O face_landmarker_v2_with_blendshapes.task -q https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task + - wget https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task -O face_landmarker_v2_with_blendshapes.task predict: "predict.py:Predictor" image: "r8.im/edenartlab/sdxl-lora-trainer" From 5f8fe5c4c834144ea404e2c8102e4af91311ff6d Mon Sep 17 00:00:00 2001 From: Gene Kogan Date: Tue, 16 Jul 2024 03:44:59 -0700 Subject: [PATCH 14/27] Update requirements.txt --- requirements.txt | 38 +++++++++++++++++++++----------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/requirements.txt b/requirements.txt index 4faf486..2845ff5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,18 +1,22 @@ -torch>=2.1.0 -torchvision>=0.16.0 -transformers>=4.38.1 -diffusers>=0.27.2 -ujson>=5.9.0 -scipy>=1.12.0 -peft>=0.10.0 -invisible-watermark>=0.2.0 +torch==2.1.0 +torchaudio==2.1.0 +torchvision==0.16.0 +transformers==4.38.0 +diffusers==0.26.0 +tokenizers==0.15.2 +huggingface-hub==0.22.2 +ujson==5.10.0 +scipy==1.14.0 +peft==0.10.0 +invisible-watermark==0.2.0 pandas==2.2.1 -numpy>=1.26.4 -opencv-python>=4.1.0.25 -mediapipe>=0.10.11 -openai>=1.14.0 -python-dotenv -prodigyopt -omegaconf -ujson -bitsandbytes \ No newline at end of file +numpy==1.26.4 +opencv-python==4.10.0.84 +mediapipe== 0.10.14 +openai==1.35.13 +python-dotenv==1.0.1 +prodigyopt==1.0 +omegaconf==2.3.0 +ujson==5.10.0 +bitsandbytes==0.43.1 +setuptools==70.3.0 From f937106f6cd869a86948c6745923e8558caac2ef Mon Sep 17 00:00:00 2001 From: aiXander Date: Tue, 16 Jul 2024 04:04:34 -0700 Subject: [PATCH 15/27] push changes --- README.md | 5 +++-- requirements.txt | 2 +- train_configs/training_args_face_sdxl.json | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 12dacdd..8aaa8c8 100755 --- a/README.md +++ b/README.md @@ -44,8 +44,9 @@ sudo curl -o /usr/local/bin/cog -L "https://github.com/replicate/cog/releases/la sudo chmod +x /usr/local/bin/cog ``` -2. Build the image with `sudo cog build` -3. Run a training run with `sudo sh cog_test_train.sh` +2. Build the image with `cog build` +3. Run a training run with `sh cog_test_train.sh` +4. You can also go into the container with `cog run /bin/bash` ## Automatic Checkpoint Evaluation diff --git a/requirements.txt b/requirements.txt index 2845ff5..ee01c4f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,7 +12,7 @@ invisible-watermark==0.2.0 pandas==2.2.1 numpy==1.26.4 opencv-python==4.10.0.84 -mediapipe== 0.10.14 +mediapipe==0.10.14 openai==1.35.13 python-dotenv==1.0.1 prodigyopt==1.0 diff --git a/train_configs/training_args_face_sdxl.json b/train_configs/training_args_face_sdxl.json index 3662f3d..a469f91 100644 --- a/train_configs/training_args_face_sdxl.json +++ b/train_configs/training_args_face_sdxl.json @@ -13,7 +13,7 @@ "ti_lr": 0.001, "ti_weight_decay": 0.0005, - "remove_ti_token_from_prompts": false, + "disable_ti": false, "text_encoder_lora_optimizer": null, "text_encoder_lora_lr": 1.0e-4, "text_encoder_lora_weight_decay": 1e-5, From fb521e9dbde0b6f26ddd11a56f8270061cc1a54b Mon Sep 17 00:00:00 2001 From: aiXander Date: Tue, 16 Jul 2024 04:39:29 -0700 Subject: [PATCH 16/27] add print flush --- cog_test_train.sh | 2 +- main.py | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/cog_test_train.sh b/cog_test_train.sh index 6e15450..9a73e8a 100644 --- a/cog_test_train.sh +++ b/cog_test_train.sh @@ -1,5 +1,5 @@ # Set GPU ID to run these jobs on: -GPU_ID="device=2" +GPU_ID="device=3" cog predict --gpus $GPU_ID \ -i name="xander_sdxl_cog" \ diff --git a/main.py b/main.py index 36406a6..1ee746f 100755 --- a/main.py +++ b/main.py @@ -197,7 +197,7 @@ def train( print(f"--- Instantaneous batch size per device = {config.train_batch_size}") print(f"--- Total batch_size (distributed + accumulation) = {total_batch_size}") print(f"--- Gradient Accumulation steps = {config.gradient_accumulation_steps}") - print(f"--- Total optimization steps = {config.max_train_steps}\n") + print(f"--- Total optimization steps = {config.max_train_steps}\n", flush = True) global_step = 0 last_save_step = 0 @@ -402,7 +402,6 @@ def train( plot_torch_hist(embedding, global_step, os.path.join(config.output_dir, 'ti_embeddings') , f"enc_{idx}_tokid_{i}: {token}", min_val=-0.05, max_val=0.05, ymax_f = 0.05, color = 'red') embedding_handler.print_token_info() - print("plotting...") if config.is_lora: # plotting this hist for full unet parameters can run OOM plot_torch_hist(unet_lora_parameters, global_step, config.output_dir, "lora_weights", min_val=-0.4, max_val=0.4, ymax_f = 0.08) plot_loss(losses, save_path=f'{config.output_dir}/losses.png') @@ -412,7 +411,6 @@ def train( plot_lrs(optimizer_collection.learning_rate_tracker, save_path=f'{config.output_dir}/learning_rates.png') plot_curve(prompt_embeds_norms, 'steps', 'norm', 'prompt_embed norms', save_path=f'{config.output_dir}/prompt_embeds_norms.png') - print("Getting ready to render validation images...") validation_prompts = render_images( pipe = pipe, render_size = config.validation_img_size, @@ -438,11 +436,11 @@ def train( if global_step % (config.max_train_steps//20) == 0: progress = (global_step / config.max_train_steps) + 0.05 print_system_info() - print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r") + print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r", flush = True) yield np.min((progress, 1.0)) if global_step > config.max_train_steps: - print("Reached max steps, stopping training!") + print("Reached max steps, stopping training!", flush = True) break # final_save @@ -473,7 +471,6 @@ def train( pretrained_model_version=config.pretrained_model["version"] ) - print("Running final inference round...") if config.debug: # Reload the entire pipe from disk + LoRa: pipe_to_use = None @@ -533,6 +530,7 @@ def train( config.job_time = time.time() - config.start_time config.training_attributes["validation_prompts"] = validation_prompts config.save_as_json(os.path.join(output_save_dir, "training_args.json")) + print("Training job complete, saving outputs...", flush = True) return config, output_save_dir From 4f084d251cab2209d84a08474e59bad6ee3a6cab Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 03:12:11 +0200 Subject: [PATCH 17/27] updates for comfyui node --- main.py | 5 ++- node.py | 76 ++++++++++++++++++++----------------------- trainer/config.py | 76 +++++++++++++++++-------------------------- trainer/models.py | 39 +++++----------------- trainer/preprocess.py | 51 +++++++++++++++++++---------- 5 files changed, 113 insertions(+), 134 deletions(-) diff --git a/main.py b/main.py index 1ee746f..9503460 100755 --- a/main.py +++ b/main.py @@ -68,7 +68,10 @@ def train( text_encoder_two, vae, unet, - ) = load_models(config.pretrained_model, config.device, weight_dtype, keep_vae_float32=0) + ), sd_model_version = load_models(config.pretrained_model, config.device, weight_dtype, keep_vae_float32=0) + + config.sd_model_version = sd_model_version + config.pretrained_model["version"] = sd_model_version # Initialize new tokens for training. embedding_handler = TokenEmbeddingsHandler( diff --git a/node.py b/node.py index d10860d..840215e 100644 --- a/node.py +++ b/node.py @@ -1,22 +1,14 @@ - import os -import shutil import tarfile import json import time -import random import torch -import numpy as np -import pandas as pd -from dotenv import load_dotenv from main import train - -from trainer.preprocess import preprocess -from trainer.models import pretrained_models -from trainer.config import TrainingConfig +from trainer.config import TrainingConfig, model_paths from trainer.utils.io import clean_filename -from trainer.utils.utils import seed_everything + +import folder_paths class Eden_LoRa_trainer: @classmethod @@ -24,9 +16,9 @@ class Eden_LoRa_trainer: return { "required": { "training_images_folder_path": ("STRING", {"default": "."}), - "lora_name": ("STRING", {"default": ""}), - "sd_model_version": (["sdxl", "sd15"], ), - "seed": ("INT", {"default": 0, "min": 0, "max": 100000}), + "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), + "lora_name": ("STRING", {"default": "Eden_LoRa"}), + "mode": (["style", "face", "object"], ), "resolution": ("INT", {"default": 512, "min": 256, "max": 768}), "train_batch_size": ("INT", {"default": 4, "min": 1, "max": 8}), "max_train_steps": ("INT", {"default": 400, "min": 50, "max": 1000}), @@ -35,17 +27,22 @@ class Eden_LoRa_trainer: "lora_rank": ("INT", {"default": 16, "min": 1, "max": 64}), "use_dora": ("BOOLEAN", {"default": False}), "n_tokens": ("INT", {"default": 2, "min": 1, "max": 3}), + "debug_mode": ("BOOLEAN", {"default": False}), + "checkpointing_steps": ("INT", {"default": 200, "min": 100, "max": 2000}), + "seed": ("INT", {"default": 0, "min": 0, "max": 100000}), } } CATEGORY = "Eden 🌱" - RETURN_TYPES = ("STRING",) + RETURN_TYPES = ("STRING", "STRING", "FLOAT") + RETURN_NAMES = ("lora_path", "embedding_path", "total_runtime_s") FUNCTION = "train_lora" - def train_lora(self, training_images_folder_path, - name = lora_name, - concept_mode = "style", - sd_model_version = "sdxl", + def train_lora(self, + training_images_folder_path, + ckpt_name, + lora_name = "eden_lora", + mode = "style", seed = 0, resolution = 521, train_batch_size = 4, @@ -54,21 +51,31 @@ class Eden_LoRa_trainer: unet_lr = 0.001, lora_rank = 16, use_dora = False, - n_tokens = 2 + n_tokens = 2, + debug_mode = False, + checkpointing_steps = 1000, ): print("Starting new training job...") + # Overwrite hardcoded paths to point to comfyUI folders: + model_paths.set_path("CLIP", os.path.join(folder_paths.models_dir, "clipseg")) + model_paths.set_path("BLIP", os.path.join(folder_paths.models_dir, "blip")) + model_paths.set_path("SR", os.path.join(folder_paths.models_dir, "upscale_models")) + model_paths.set_path("SD", os.path.join(folder_paths.models_dir, "checkpoints")) + + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + config = TrainingConfig( - name="test", + name=lora_name, lora_training_urls=training_images_folder_path, - concept_mode=concept_mode, - sd_model_version=sd_model_version, + concept_mode=mode, + ckpt_path=ckpt_path, seed=seed, resolution=resolution, train_batch_size=train_batch_size, max_train_steps=max_train_steps, - checkpointing_steps=10000, + checkpointing_steps=checkpointing_steps, ti_lr=ti_lr, unet_lr=unet_lr, lora_rank=lora_rank, @@ -76,7 +83,7 @@ class Eden_LoRa_trainer: caption_model="blip", n_tokens=n_tokens, verbose=True, - debug=True, + debug=debug_mode, ) with torch.inference_mode(False): @@ -89,22 +96,8 @@ class Eden_LoRa_trainer: break validation_grid_img_path = os.path.join(output_save_dir, "validation_grid.jpg") - out_path = f"{clean_filename(lora_name)}_eden_concept_lora_{int(time.time())}.tar" directory = cogPath(output_save_dir) - with tarfile.open(out_path, "w") as tar: - print("Adding files to tar...") - for file_path in directory.rglob("*"): - print(file_path) - arcname = file_path.relative_to(directory) - tar.add(file_path, arcname=arcname) - - # Add instructions README: - tar.add("instructions_README.md", arcname="README.md") - tar.add("comfyUI_workflow_lora_txt2img.json", arcname="comfyUI_workflow_lora_txt2img.json") - if sd_model_version == "sd15": - tar.add("comfyUI_workflow_lora_adiff.json", arcname="comfyUI_workflow_lora_adiff.json") - attributes = {} attributes['grid_prompts'] = config.training_attributes["validation_prompts"] attributes['job_time_seconds'] = config.job_time @@ -112,4 +105,7 @@ class Eden_LoRa_trainer: print(f"LORA training finished in {config.job_time:.1f} seconds") print(f"Returning {out_path}") - return (out_path,) \ No newline at end of file + lora_path = os.path.join(output_save_dir, config.name + ".safetensors") + embedding_path = os.path.join(output_save_dir, config.name + "_embeddings.safetensors") + + return (lora_path, embedding_path, config.job_time) \ No newline at end of file diff --git a/trainer/config.py b/trainer/config.py index c8b3e8f..16480d4 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -3,62 +3,42 @@ from datetime import datetime from pydantic import BaseModel import json, time, os from typing import Literal -from trainer.models import pretrained_models from trainer.utils.utils import pick_best_gpu_id -import subprocess -def run_command(command): - try: - result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True) - return result.stdout - except subprocess.CalledProcessError as e: - return f"Error: {e.stderr}" +class ModelPaths: + def __init__(self): + self.paths = { + "BLIP": "./cache", + "CLIP": "./cache", + "SR": "./cache", + "SD": "./models", + } + def get_path(self, key): + return self.paths.get(key, None) -print("=================================================================") -print("=================================================================") + def set_path(self, key, path): + if key in self.paths: + self.paths[key] = path -print("NVIDIA Driver Version:") -nvidia_smi_output = run_command("nvidia-smi --query-gpu=driver_version --format=csv,noheader") -print(nvidia_smi_output.strip()) # strip() removes any trailing newlines +model_paths = ModelPaths() -print("\nNVCC Version:") -print(run_command("nvcc --version")) +# Default download urls in case no local model is found: +SDXL_URL = "https://huggingface.co/RunDiffusion/Juggernaut-XL-v6/resolve/main/juggernautXL_version6Rundiffusion.safetensors" +SD15_URL = "https://huggingface.co/KamCastle/jugg/resolve/main/juggernaut_reborn.safetensors" -print("\nCUDNN Libraries:") -print(run_command("ldconfig -p | grep cudnn")) - -print("\nFull NVIDIA-SMI Output:") -print(run_command("nvidia-smi")) - -# Run nvcc --version -nvcc_version = run_command("nvcc --version") -print("NVCC Version:") -print(nvcc_version) - -# Run ldconfig -p | grep cudnn -cudnn_info = run_command("ldconfig -p | grep cudnn") -print("\nCUDNN Libraries:") -print(cudnn_info) - -# Check if CUDA is available -import torch -if torch.cuda.is_available(): - print(f"Torch version: {torch.__version__}") - print(f"CUDA Version: {torch.version.cuda}") - print(f"cuDNN Version: {torch.backends.cudnn.version()}") -else: - print("CUDA is not available.") - -print("-----------------------------------------------------------------") -print("-----------------------------------------------------------------") +pretrained_models = { + "sdxl": {"path": os.path.join(model_paths.get_path("SD"), os.path.basename(SDXL_URL)), "url": SDXL_URL, "version": "sdxl"}, + "sd15": {"path": os.path.join(model_paths.get_path("SD"), os.path.basename(SD15_URL)), "url": SD15_URL, "version": "sd15"} +} class TrainingConfig(BaseModel): lora_training_urls: str concept_mode: Literal["face", "style", "object"] caption_prefix: str = "" # hardcoding this will inject TOK manually and skip the chatgpt token injection step, not recommended unless you know what you're doing caption_model: Literal["gpt4-v", "blip"] = "blip" - sd_model_version: Literal["sdxl", "sd15"] + sd_model_version: Literal["sdxl", "sd15", None] = None + ckpt_path: str = None # optional hardcoded checkpoint path pretrained_model: dict = None seed: Union[int, None] = None resolution: int = 512 @@ -106,7 +86,7 @@ class TrainingConfig(BaseModel): clipseg_temperature: float = 0.5 # temperature for the CLIPSeg mask n_sample_imgs: int = 4 name: str = None - output_dir: str = "lora_models/unnamed" + output_dir: str = "eden_lora_training_runs" debug: bool = False allow_tf32: bool = True disable_ti: bool = False @@ -140,7 +120,11 @@ class TrainingConfig(BaseModel): def __init__(self, **data): super().__init__(**data) - self.pretrained_model = pretrained_models[self.sd_model_version] + + if not self.ckpt_path: + self.pretrained_model = pretrained_models[self.sd_model_version] + else: + self.pretrained_model = {"path": self.ckpt_path, "url": None, "version": None} # add some metrics to the foldername: lora_str = "dora" if self.use_dora else "lora" @@ -149,7 +133,7 @@ class TrainingConfig(BaseModel): if not self.name: self.name = f"{os.path.basename(self.output_dir)}_{self.concept_mode}_{lora_str}_{self.sd_model_version}_{timestamp_short}" - self.output_dir = self.output_dir + f"--{timestamp_short}-{self.sd_model_version}_{self.concept_mode}_{lora_str}_{self.resolution}_{self.prodigy_d_coef}_{self.caption_model}_{self.max_train_steps}" + self.output_dir = self.output_dir + f"/{self.name}/" + f"{timestamp_short}-{self.sd_model_version}_{self.concept_mode}_{lora_str}_{self.resolution}_{self.prodigy_d_coef}_{self.caption_model}_{self.max_train_steps}" os.makedirs(self.output_dir, exist_ok=True) if self.seed is None: diff --git a/trainer/models.py b/trainer/models.py index 959961d..04cc21d 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -4,44 +4,23 @@ import subprocess import torch from diffusers import AutoencoderKL, DDPMScheduler, EulerDiscreteScheduler, UNet2DConditionModel, StableDiffusionPipeline, StableDiffusionXLPipeline -############################################################################################################ - -SDXL_MODEL_CACHE = "./models/juggernaut_v6.safetensors" -SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernautXL_v6.safetensors" - -#SDXL_MODEL_CACHE = "./models/sd_xl_base_1.0_0.9vae.safetensors" -#SDXL_URL = "https://edenartlab-lfs.s3.amazonaws.com/comfyui/models2/checkpoints/sd_xl_base_1.0_0.9vae.safetensors" - -SD15_MODEL_CACHE = "./models/juggernaut_reborn.safetensors" -SD15_URL = "https://edenartlab-lfs.s3.amazonaws.com/models/checkpoints/juggernaut_reborn.safetensors" - -#SD15_MODEL_CACHE = "./models/photon_v1.safetensors" -#SD15_URL = "https://civitai.com/api/download/models/90072" - -pretrained_models = { - "sdxl": {"path": SDXL_MODEL_CACHE, "url": SDXL_URL, "version": "sdxl"}, - "sd15": {"path": SD15_MODEL_CACHE, "url": SD15_URL, "version": "sd15"} -} - -############################################################################################################ - - def load_models(pretrained_model, device, weight_dtype = torch.float16, keep_vae_float32 = False): - if not isinstance(pretrained_model, dict) or 'path' not in pretrained_model or 'version' not in pretrained_model: - raise ValueError("pretrained_model must be a dict with 'path' and 'version' keys") - # check if the model is already downloaded: if not os.path.exists(pretrained_model['path']): download_weights(pretrained_model['url'], pretrained_model['path']) print(f"Loading model weights from {pretrained_model['path']} with dtype: {weight_dtype}...") - if pretrained_model['version'] == "sd15": - pipe = StableDiffusionPipeline.from_single_file( - pretrained_model['path'], torch_dtype=weight_dtype, use_safetensors=True) - else: + try: pipe = StableDiffusionXLPipeline.from_single_file( pretrained_model['path'], torch_dtype=weight_dtype, use_safetensors=True) + sd_model_version = "sdxl" + except: + pipe = StableDiffusionPipeline.from_single_file( + pretrained_model['path'], torch_dtype=weight_dtype, use_safetensors=True) + sd_model_version = "sd15" + + print(f"Loaded {sd_model_version} model!") pipe = pipe.to(device, dtype=weight_dtype) noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config) @@ -79,7 +58,7 @@ def load_models(pretrained_model, device, weight_dtype = torch.float16, keep_vae text_encoder_two, vae, unet, - ) + ), sd_model_version def download_weights(url, dest): start = time.time() diff --git a/trainer/preprocess.py b/trainer/preprocess.py index 74d61f5..564f497 100755 --- a/trainer/preprocess.py +++ b/trainer/preprocess.py @@ -1,7 +1,3 @@ -# Have SwinIR upsample -# Have BLIP auto caption -# Have CLIPSeg auto mask concept - import gc import fnmatch import mimetypes @@ -25,6 +21,7 @@ import numpy as np import pandas as pd import torch from tqdm import tqdm + from transformers import ( BlipForConditionalGeneration, Blip2ForConditionalGeneration, @@ -38,13 +35,13 @@ from transformers import ( from trainer.utils.io import download_and_prep_training_data from trainer.utils.utils import fix_prompt +from trainer.config import model_paths import re import openai from openai import OpenAI from dotenv import load_dotenv load_dotenv() - try: OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") client = OpenAI(api_key=OPENAI_API_KEY) @@ -54,8 +51,6 @@ except: client = None print("WARNING: Could not find OPENAI_API_KEY in .env, disabling gpt prompt generation.") -MODEL_PATH = "./cache" - # Put some boundaries to make the gpt pass work well: (very long text often confuses the model and also costs more money...) MIN_GPT_PROMPTS = 3 MAX_GPT_PROMPTS = 50 @@ -139,7 +134,7 @@ def swin_ir_sr( """ model = Swin2SRForImageSuperResolution.from_pretrained( - model_id, cache_dir=MODEL_PATH + model_id, cache_dir = model_paths.get_path("SR") ).to(device) processor = Swin2SRImageProcessor() @@ -193,9 +188,9 @@ def clipseg_mask_generator( model = None if any(target_prompts): - processor = CLIPSegProcessor.from_pretrained(model_id, cache_dir=MODEL_PATH) + processor = CLIPSegProcessor.from_pretrained(model_id, cache_dir = model_paths.get_path("CLIP")) model = CLIPSegForImageSegmentation.from_pretrained( - model_id, cache_dir=MODEL_PATH + model_id, cache_dir = model_paths.get_path("CLIP") ).to(device) masks = [] @@ -406,16 +401,16 @@ def blip_caption_dataset( print(f"Using model {model_id} for image captioning...") device=torch.device("cuda" if torch.cuda.is_available() else "cpu") - + if "blip2" in model_id: - processor = Blip2Processor.from_pretrained(model_id, cache_dir=MODEL_PATH) + processor = Blip2Processor.from_pretrained(model_id, cache_dir = model_paths.get_path("BLIP")) model = Blip2ForConditionalGeneration.from_pretrained( - model_id, cache_dir=MODEL_PATH, torch_dtype=torch.float16 + model_id, cache_dir = model_paths.get_path("BLIP"), torch_dtype=torch.float16 ).to(device) else: - processor = BlipProcessor.from_pretrained(model_id, cache_dir=MODEL_PATH) + processor = BlipProcessor.from_pretrained(model_id, cache_dir = model_paths.get_path("BLIP")) model = BlipForConditionalGeneration.from_pretrained( - model_id, cache_dir=MODEL_PATH, torch_dtype=torch.float16 + model_id, cache_dir = model_paths.get_path("BLIP"), torch_dtype=torch.float16 ).to(device) for i, image in enumerate(tqdm(images)): @@ -643,6 +638,30 @@ def augment_image(image): def round_to_nearest_multiple(x, multiple): return int(float(multiple) * round(float(x) / float(multiple))) +''' +For Stable Diffusion 1.5, outputs are optimised around 512x512 pixels. Many common fine-tuned versions of SD1.5 are optimised around 768x768. The best resolutions for common aspect ratios are typically: +1:1 (square): 512x512, 768x768 +3:2 (landscape): 768x512 +2:3 (portrait): 512x768 +4:3 (landscape): 768x576 +3:4 (portrait): 576x768 +16:9 (widescreen): 912x512 +9:16 (tall): 512x912 + +For SDXL, outputs are optimised around 1024x1024 pixels. The best resolutions for common aspect ratios are typically: +stable-diffusion-xl-1024-v0-9 supports generating images at the following dimensions: +1024 x 1024 +1152 x 896 +896 x 1152 +1216 x 832 +832 x 1216 +1344 x 768 +768 x 1344 +1536 x 640 +640 x 1536 + +''' + def calculate_new_dimensions(target_size, target_aspect_ratio): """ Calculate the new width and height given a target size and aspect ratio. @@ -661,8 +680,6 @@ def calculate_new_dimensions(target_size, target_aspect_ratio): return [new_width, new_height] - - def load_and_save_masks_and_captions( config, concept_mode: str, From d8ee9ee522720f16fdffe8b8eb36dcc0679f0997 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 03:14:31 +0200 Subject: [PATCH 18/27] use gpt-4o --- trainer/preprocess.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trainer/preprocess.py b/trainer/preprocess.py index 564f497..e79ca6e 100755 --- a/trainer/preprocess.py +++ b/trainer/preprocess.py @@ -401,7 +401,7 @@ def blip_caption_dataset( print(f"Using model {model_id} for image captioning...") device=torch.device("cuda" if torch.cuda.is_available() else "cpu") - + if "blip2" in model_id: processor = Blip2Processor.from_pretrained(model_id, cache_dir = model_paths.get_path("BLIP")) model = Blip2ForConditionalGeneration.from_pretrained( @@ -468,7 +468,7 @@ def gpt4_v_get_description(config, images): base64_image = prep_img_for_gpt_api(img, max_size=(1024, 1024)) payload = { - "model": "gpt-4-turbo", + "model": "gpt-4o", "messages": [ { "role": "user", From 2a0761ae75e8035dfef4bd6ec3b1ca90cef7d4db Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 03:16:41 +0200 Subject: [PATCH 19/27] use gpt-4o --- trainer/preprocess.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trainer/preprocess.py b/trainer/preprocess.py index 74d61f5..4f9b51d 100755 --- a/trainer/preprocess.py +++ b/trainer/preprocess.py @@ -510,7 +510,7 @@ def gpt4_v_caption_dataset( base64_image = prep_img_for_gpt_api(img, max_size=(512, 512)) payload = { - "model": "gpt-4-turbo", + "model": "gpt-4o", "messages": [ { "role": "user", From af16d1edcc9235c245089b5999d3a85f675d83cf Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 03:26:51 +0200 Subject: [PATCH 20/27] update configs --- scripts/create_hyperparam_sweep.py | 17 +++++++----- train_configs/training_args_beeple.json | 30 --------------------- train_configs/training_args_face_sd15.json | 2 +- train_configs/training_args_face_sdxl.json | 2 +- train_configs/training_args_object.json | 4 +-- train_configs/training_args_style_sd15.json | 2 +- train_configs/training_args_style_sdxl.json | 5 ++-- trainer/config.py | 3 ++- 8 files changed, 21 insertions(+), 44 deletions(-) delete mode 100644 train_configs/training_args_beeple.json diff --git a/scripts/create_hyperparam_sweep.py b/scripts/create_hyperparam_sweep.py index 552fe74..0e2d1f3 100644 --- a/scripts/create_hyperparam_sweep.py +++ b/scripts/create_hyperparam_sweep.py @@ -35,12 +35,12 @@ def hamming_distance(dict1, dict2): ####################################################################################### # Setup the base experiment config: -exp_name = "grimes" +exp_name = "beeple" caption_prefix = "" mask_target_prompts = "" n_exp = 200 # how many random experiment settings to generate -min_hamming_distance = 3 # min_n_params that have to be different from any previous experiment to be scheduled - +min_hamming_distance = 1 # min_n_params that have to be different from any previous experiment to be scheduled +nohup = True output_sh_path = f"gridsearch_configs/{exp_name}.sh" # Define training hyperparameters and their possible values @@ -55,7 +55,7 @@ hyperparameters = { ], "concept_mode": ['style'], - "sample_imgs_lora_scale": [0.9], + "sample_imgs_lora_scale": [0.8], "disable_ti": ['false', 'true'], "seed": [0], "resolution": [512], @@ -72,7 +72,7 @@ hyperparameters = { "token_warmup_steps": [0], "tok_cov_reg_w": [2000], - "unet_lr": [0.001, 0.0001], + "unet_lr": [0.0002, 0.00005], "lora_alpha_multiplier": [1.0], "prodigy_d_coef": [1.0], "lora_weight_decay": [0.001], @@ -149,7 +149,12 @@ def generate_sh_script(folder_path, output_sh_path): # Write a command for each JSON file for json_file in json_files: - command = f"python main.py {os.path.join(folder_path, json_file)}\n" + file_path = os.path.join("scripts/", folder_path, json_file) + command = f"python main.py {file_path}\n" + + if nohup: + command = f"nohup {command} > {file_path.replace('.json', '.log')} 2>&1 &\n" + sh_file.write(command) generate_sh_script(config_output_dir, output_sh_path) diff --git a/train_configs/training_args_beeple.json b/train_configs/training_args_beeple.json deleted file mode 100644 index e53dcad..0000000 --- a/train_configs/training_args_beeple.json +++ /dev/null @@ -1,30 +0,0 @@ -{ - "output_dir": "lora_models/beeple", - "sd_model_version": "sdxl", - "lora_training_urls": "/home/rednax/SSD2TB/Github_repos/Eden/images/beeple_large", - "concept_mode": "style", - "sample_imgs_lora_scale": 0.9, - "seed": 0, - "resolution": 512, - "train_batch_size": 4, - "n_sample_imgs": 8, - "max_train_steps": 1200, - "token_warmup_steps": 0, - "checkpointing_steps": 200, - "ti_lr": 0.001, - "ti_weight_decay": 0.0005, - - "disable_ti": true, - "text_encoder_lora_optimizer": null, - "text_encoder_lora_lr": 1.0e-4, - "text_encoder_lora_weight_decay": 1e-5, - "text_encoder_lora_rank": 12, - - "unet_lr": 0.0001, - "lora_rank": 4, - "use_dora": false, - "is_lora": false, - "unet_optimizer_type": "AdamW8bit", - "caption_model": "blip", - "debug": true -} \ No newline at end of file diff --git a/train_configs/training_args_face_sd15.json b/train_configs/training_args_face_sd15.json index dc7f423..8ed7052 100644 --- a/train_configs/training_args_face_sd15.json +++ b/train_configs/training_args_face_sd15.json @@ -1,5 +1,5 @@ { - "output_dir": "lora_models/test_xander", + "name": "xander_sd15", "sd_model_version": "sd15", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", "concept_mode": "face", diff --git a/train_configs/training_args_face_sdxl.json b/train_configs/training_args_face_sdxl.json index a469f91..7552ef7 100644 --- a/train_configs/training_args_face_sdxl.json +++ b/train_configs/training_args_face_sdxl.json @@ -1,5 +1,5 @@ { - "output_dir": "lora_models/test_xander", + "name": "xander_sdxl", "sd_model_version": "sdxl", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", "concept_mode": "face", diff --git a/train_configs/training_args_object.json b/train_configs/training_args_object.json index 4334731..ed07af3 100644 --- a/train_configs/training_args_object.json +++ b/train_configs/training_args_object.json @@ -1,7 +1,7 @@ { - "output_dir": "lora_models/test_xander", + "name": "banny_sd15", "sd_model_version": "sd15", - "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander.zip", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/banny.zip", "concept_mode": "face", "sample_imgs_lora_scale": 0.8, "seed": 0, diff --git a/train_configs/training_args_style_sd15.json b/train_configs/training_args_style_sd15.json index 10ad4d8..75b1c5c 100644 --- a/train_configs/training_args_style_sd15.json +++ b/train_configs/training_args_style_sd15.json @@ -1,5 +1,5 @@ { - "output_dir": "lora_models/test_clipx", + "name": "clipx_sd15", "sd_model_version": "sd15", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/clipx_tiny.zip", "concept_mode": "style", diff --git a/train_configs/training_args_style_sdxl.json b/train_configs/training_args_style_sdxl.json index e4ca178..2dc9884 100644 --- a/train_configs/training_args_style_sdxl.json +++ b/train_configs/training_args_style_sdxl.json @@ -1,9 +1,10 @@ { - "output_dir": "lora_models/test_clipx", + "name": "clipx_sdxl", "sd_model_version": "sdxl", "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/clipx_tiny.zip", "concept_mode": "style", - "seed": 0, + "sample_imgs_lora_scale": 0.7, + "seed": 1, "resolution": 512, "train_batch_size": 4, "n_sample_imgs": 6, diff --git a/trainer/config.py b/trainer/config.py index 16480d4..1f16121 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -24,7 +24,8 @@ class ModelPaths: model_paths = ModelPaths() # Default download urls in case no local model is found: -SDXL_URL = "https://huggingface.co/RunDiffusion/Juggernaut-XL-v6/resolve/main/juggernautXL_version6Rundiffusion.safetensors" +#SDXL_URL = "https://huggingface.co/RunDiffusion/Juggernaut-XL-v6/resolve/main/juggernautXL_version6Rundiffusion.safetensors" +SDXL_URL = "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0_0.9vae.safetensors" SD15_URL = "https://huggingface.co/KamCastle/jugg/resolve/main/juggernaut_reborn.safetensors" pretrained_models = { From 85599dc72537cc2b5c8beb8a51270b68279ab88f Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 03:43:43 +0200 Subject: [PATCH 21/27] minor changes --- main.py | 37 ++++++++++++++++++------------------- node.py | 2 +- trainer/config.py | 2 +- trainer/models.py | 2 +- 4 files changed, 21 insertions(+), 22 deletions(-) diff --git a/main.py b/main.py index 9503460..fb212a7 100755 --- a/main.py +++ b/main.py @@ -25,6 +25,7 @@ from trainer.loss import compute_diffusion_loss, compute_grad_norm, Conditioning from trainer.inference import render_images, get_conditioning_signals from trainer.preprocess import preprocess from trainer.utils.io import make_validation_img_grid + from trainer.optimizer import ( OptimizerCollection, get_optimizer_and_peft_models_text_encoder_lora, @@ -33,10 +34,24 @@ from trainer.optimizer import ( get_unet_optimizer ) -def train( - config: TrainingConfig, -): +def train(config: TrainingConfig): + seed_everything(config.seed) + weight_dtype = dtype_map[config.weight_type] + + ( + pipe, + tokenizer_one, + tokenizer_two, + noise_scheduler, + text_encoder_one, + text_encoder_two, + vae, + unet, + ), sd_model_version = load_models(config.pretrained_model, config.device, weight_dtype) + + config.sd_model_version = sd_model_version + config.pretrained_model["version"] = sd_model_version config, input_dir = preprocess( config, @@ -57,22 +72,6 @@ def train( if config.allow_tf32: torch.backends.cuda.matmul.allow_tf32 = True - weight_dtype = dtype_map[config.weight_type] - - ( - pipe, - tokenizer_one, - tokenizer_two, - noise_scheduler, - text_encoder_one, - text_encoder_two, - vae, - unet, - ), sd_model_version = load_models(config.pretrained_model, config.device, weight_dtype, keep_vae_float32=0) - - config.sd_model_version = sd_model_version - config.pretrained_model["version"] = sd_model_version - # Initialize new tokens for training. embedding_handler = TokenEmbeddingsHandler( text_encoders = [text_encoder_one, text_encoder_two], diff --git a/node.py b/node.py index 840215e..bf31d37 100644 --- a/node.py +++ b/node.py @@ -28,7 +28,7 @@ class Eden_LoRa_trainer: "use_dora": ("BOOLEAN", {"default": False}), "n_tokens": ("INT", {"default": 2, "min": 1, "max": 3}), "debug_mode": ("BOOLEAN", {"default": False}), - "checkpointing_steps": ("INT", {"default": 200, "min": 100, "max": 2000}), + "checkpointing_steps": ("INT", {"default": 200, "min": 10, "max": 2000}), "seed": ("INT", {"default": 0, "min": 0, "max": 100000}), } } diff --git a/trainer/config.py b/trainer/config.py index 1f16121..0b1fdd8 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -134,7 +134,7 @@ class TrainingConfig(BaseModel): if not self.name: self.name = f"{os.path.basename(self.output_dir)}_{self.concept_mode}_{lora_str}_{self.sd_model_version}_{timestamp_short}" - self.output_dir = self.output_dir + f"/{self.name}/" + f"{timestamp_short}-{self.sd_model_version}_{self.concept_mode}_{lora_str}_{self.resolution}_{self.prodigy_d_coef}_{self.caption_model}_{self.max_train_steps}" + self.output_dir = self.output_dir + f"/{self.name}/" + f"{timestamp_short}-{self.concept_mode}_{lora_str}_{self.resolution}_{self.prodigy_d_coef}_{self.caption_model}_{self.max_train_steps}" os.makedirs(self.output_dir, exist_ok=True) if self.seed is None: diff --git a/trainer/models.py b/trainer/models.py index 04cc21d..4b3467f 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -36,7 +36,7 @@ def load_models(pretrained_model, device, weight_dtype = torch.float16, keep_vae else: vae.to(device, dtype=weight_dtype) if weight_dtype != torch.float32: - print(f"Warning: VAE will be loaded as {weight_dtype}, this is fine for inference but might not be for training..") + print(f"Warning: VAE will be loaded as {weight_dtype}, this is fine for inference but may not be ideal for training..?") unet.to(device, dtype=weight_dtype) text_encoder_one.requires_grad_(False) From c891203365e4bffff1c91b5802caa18fd2a75bc6 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 04:00:08 +0200 Subject: [PATCH 22/27] update gitignore --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 1e034e9..8f761e9 100644 --- a/.gitignore +++ b/.gitignore @@ -4,13 +4,13 @@ __pycache__ models lora_models* +eden_lora_training_runs/ datasets *.tar .env .cog .huggingface -train.py rendered_images* gridsearch* From 01523350e1d3ba231f17884ea5c97b1c9f0a1a1f Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 04:00:21 +0200 Subject: [PATCH 23/27] cleanup save dir --- main.py | 10 ++-------- trainer/checkpoint.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/main.py b/main.py index fb212a7..0b2af67 100755 --- a/main.py +++ b/main.py @@ -437,7 +437,7 @@ def train(config: TrainingConfig): if global_step % (config.max_train_steps//20) == 0: progress = (global_step / config.max_train_steps) + 0.05 - print_system_info() + #print_system_info() print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r", flush = True) yield np.min((progress, 1.0)) @@ -512,13 +512,6 @@ def train(config: TrainingConfig): img_grid_path = make_validation_img_grid(output_save_dir) shutil.copy(img_grid_path, os.path.join(os.path.dirname(output_save_dir), f"validation_grid_{global_step:04d}.jpg")) - # Remove unneeded checkpoints if they exist in the output directory: - to_remove = ["pytorch_lora_weights.safetensors", "adapter_model.safetensors"] - for file in to_remove: - file_path = os.path.join(output_save_dir, file) - if os.path.exists(file_path): - os.remove(file_path) - else: print(f"Skipping final save, {output_save_dir} already exists") @@ -546,6 +539,7 @@ if __name__ == "__main__": print("Starting new LoRa training run with config:") print(config) + print("------------------------------------------") for progress in train(config=config): print(f"Progress: {(100*progress):.2f}%", end="\r") diff --git a/trainer/checkpoint.py b/trainer/checkpoint.py index 1789287..7018be4 100644 --- a/trainer/checkpoint.py +++ b/trainer/checkpoint.py @@ -184,12 +184,21 @@ def save_checkpoint( convert_pytorch_lora_safetensors_to_webui( pytorch_lora_weights_filename=os.path.join(output_dir, "pytorch_lora_weights.safetensors"), - output_filename=os.path.join(output_dir, f"{name}.safetensors") + output_filename=os.path.join(output_dir, f"{name}_{pretrained_model_version}.safetensors") ) else: # Save the entire, finetuned unet weights: unet.save_pretrained(save_directory = output_dir) + # Remove unneeded checkpoints if they exist in the output directory: TODO clean this up so they are never needed in the first place.. + to_remove = ["pytorch_lora_weights.safetensors", "adapter_model.safetensors"] + for file in to_remove: + file_path = os.path.join(output_dir, file) + if os.path.exists(file_path): + os.remove(file_path) + + return + def load_checkpoint( pretrained_model_version: str, pretrained_model_path: str, From 598c4204be1ba21085843ed1108b87d7d3b6a5d2 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 04:05:10 +0200 Subject: [PATCH 24/27] add test config --- train_configs/test.json | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 train_configs/test.json diff --git a/train_configs/test.json b/train_configs/test.json new file mode 100644 index 0000000..55d83d1 --- /dev/null +++ b/train_configs/test.json @@ -0,0 +1,27 @@ +{ + "name": "xander_test", + "sd_model_version": "sdxl", + "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander_5.zip", + "concept_mode": "face", + "seed": 1, + "resolution": 512, + "train_batch_size": 4, + "n_sample_imgs": 4, + "max_train_steps": 200, + "token_warmup_steps": 0, + "checkpointing_steps": 100, + "ti_lr": 0.001, + "ti_weight_decay": 0.0005, + + "disable_ti": false, + "text_encoder_lora_optimizer": null, + "text_encoder_lora_lr": 1.0e-4, + "text_encoder_lora_weight_decay": 1e-5, + "text_encoder_lora_rank": 12, + + "unet_lr": 0.001, + "lora_rank": 16, + "use_dora": false, + "caption_model": "blip", + "debug": true +} \ No newline at end of file From ee91f8c0c7ba7c5b22e51f36550870150c655677 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 04:07:00 +0200 Subject: [PATCH 25/27] update naming conventions --- trainer/checkpoint.py | 4 ++-- trainer/models.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/trainer/checkpoint.py b/trainer/checkpoint.py index 7018be4..8e8d524 100644 --- a/trainer/checkpoint.py +++ b/trainer/checkpoint.py @@ -135,7 +135,7 @@ def save_checkpoint( embedding_handler.save_embeddings( os.path.join( output_dir, - f"{name}_embeddings.safetensors" + f"{name}_{pretrained_model_version}_embeddings.safetensors" ) ) @@ -184,7 +184,7 @@ def save_checkpoint( convert_pytorch_lora_safetensors_to_webui( pytorch_lora_weights_filename=os.path.join(output_dir, "pytorch_lora_weights.safetensors"), - output_filename=os.path.join(output_dir, f"{name}_{pretrained_model_version}.safetensors") + output_filename=os.path.join(output_dir, f"{name}_{pretrained_model_version}_LoRa.safetensors") ) else: # Save the entire, finetuned unet weights: diff --git a/trainer/models.py b/trainer/models.py index 4b3467f..66e8a53 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -9,7 +9,7 @@ def load_models(pretrained_model, device, weight_dtype = torch.float16, keep_vae if not os.path.exists(pretrained_model['path']): download_weights(pretrained_model['url'], pretrained_model['path']) - print(f"Loading model weights from {pretrained_model['path']} with dtype: {weight_dtype}...") + print(f"Loading model weights from {os.path.abspath(pretrained_model['path'])} with dtype: {weight_dtype}...") try: pipe = StableDiffusionXLPipeline.from_single_file( From 3cd17086f3a11e3634a2d4e98f391749558f7cb9 Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 05:50:04 +0200 Subject: [PATCH 26/27] trainer comfyui v1 --- main.py | 3 ++- node.py | 23 +++++++++++++++++------ trainer/models.py | 2 +- 3 files changed, 20 insertions(+), 8 deletions(-) diff --git a/main.py b/main.py index 0b2af67..3cbcbe7 100755 --- a/main.py +++ b/main.py @@ -473,7 +473,7 @@ def train(config: TrainingConfig): pretrained_model_version=config.pretrained_model["version"] ) - if config.debug: + if config.debug and 0: # Reload the entire pipe from disk + LoRa: pipe_to_use = None checkpoint_folder = output_save_dir @@ -526,6 +526,7 @@ def train(config: TrainingConfig): config.training_attributes["validation_prompts"] = validation_prompts config.save_as_json(os.path.join(output_save_dir, "training_args.json")) print("Training job complete, saving outputs...", flush = True) + print("------------------------------------------") return config, output_save_dir diff --git a/node.py b/node.py index bf31d37..e212e91 100644 --- a/node.py +++ b/node.py @@ -3,6 +3,8 @@ import tarfile import json import time import torch +import numpy as np +from PIL import Image from main import train from trainer.config import TrainingConfig, model_paths @@ -34,8 +36,8 @@ class Eden_LoRa_trainer: } CATEGORY = "Eden 🌱" - RETURN_TYPES = ("STRING", "STRING", "FLOAT") - RETURN_NAMES = ("lora_path", "embedding_path", "total_runtime_s") + RETURN_TYPES = ("IMAGE", "STRING", "STRING", "STRING") + RETURN_NAMES = ("sample_images", "lora_path", "embedding_path", "total_runtime_s") FUNCTION = "train_lora" def train_lora(self, @@ -91,21 +93,30 @@ class Eden_LoRa_trainer: while True: try: progress_f = next(train_generator) + print("---- Progress ----") + print(progress_f) except StopIteration as e: config, output_save_dir = e.value # Capture the return value + print("---- Training finished at StopIteration! ----") break validation_grid_img_path = os.path.join(output_save_dir, "validation_grid.jpg") - directory = cogPath(output_save_dir) attributes = {} attributes['grid_prompts'] = config.training_attributes["validation_prompts"] attributes['job_time_seconds'] = config.job_time - print(f"LORA training finished in {config.job_time:.1f} seconds") - print(f"Returning {out_path}") + print(f"LORA training node finished in {config.job_time:.1f} seconds") + print("---------- Made by Eden.art 🌱 ----------") lora_path = os.path.join(output_save_dir, config.name + ".safetensors") embedding_path = os.path.join(output_save_dir, config.name + "_embeddings.safetensors") - return (lora_path, embedding_path, config.job_time) \ No newline at end of file + # Load the grid image: + grid_image = Image.open(validation_grid_img_path) + grid_image = np.array(grid_image).astype(np.float32) / 255.0 + grid_image = torch.from_numpy(grid_image)[None,] + + final_msg = f"LoRa trained in {config.job_time/60:.1f} minutes. Files saved at {output_save_dir}" + + return (grid_image, lora_path, embedding_path, final_msg) \ No newline at end of file diff --git a/trainer/models.py b/trainer/models.py index 66e8a53..7bc90b8 100644 --- a/trainer/models.py +++ b/trainer/models.py @@ -43,7 +43,7 @@ def load_models(pretrained_model, device, weight_dtype = torch.float16, keep_vae text_encoder_one.to(device, dtype=weight_dtype) tokenizer_two = text_encoder_two = None - if pretrained_model['version'] == "sdxl": + if sd_model_version == "sdxl": tokenizer_two = pipe.tokenizer_2 text_encoder_two = pipe.text_encoder_2 text_encoder_two.requires_grad_(False) From 6a4fa1ef8ddf7f8be4b04b0fef5bc4c6c4d9c56f Mon Sep 17 00:00:00 2001 From: xander Date: Thu, 18 Jul 2024 21:04:43 +0200 Subject: [PATCH 27/27] add progressbar to node --- main.py | 49 +++++++++++++++++++++++++------------------------ node.py | 24 ++++++++++++++++-------- 2 files changed, 41 insertions(+), 32 deletions(-) diff --git a/main.py b/main.py index 3cbcbe7..f9219d2 100755 --- a/main.py +++ b/main.py @@ -370,7 +370,7 @@ def train(config: TrainingConfig): token_stds[f'text_encoder_{idx}'][std_i].append(embedding_stds[std_i].item()) # Print some statistics: - if config.debug and (global_step % config.checkpointing_steps == 0) and (global_step < (config.max_train_steps - 25)) and global_step > 0: + if (global_step % config.checkpointing_steps == 0) and (global_step < (config.max_train_steps - 25)) and global_step > 0: output_save_dir = f"{checkpoint_dir}/checkpoint-{global_step}" os.makedirs(output_save_dir, exist_ok=True) @@ -391,28 +391,29 @@ def train(config: TrainingConfig): ) last_save_step = global_step - token_embeddings, trainable_tokens = embedding_handler.get_trainable_embeddings() - for idx, text_encoder in enumerate(text_encoders): - if text_encoder is None: - continue - n = len(token_embeddings[f'txt_encoder_{idx}']) - for i in range(n): - token = trainable_tokens[f'txt_encoder_{idx}'][i] - # Strip any backslashes from the token name: - token = token.replace("/", "_") - embedding = token_embeddings[f'txt_encoder_{idx}'][i] - plot_torch_hist(embedding, global_step, os.path.join(config.output_dir, 'ti_embeddings') , f"enc_{idx}_tokid_{i}: {token}", min_val=-0.05, max_val=0.05, ymax_f = 0.05, color = 'red') + if config.debug: + token_embeddings, trainable_tokens = embedding_handler.get_trainable_embeddings() + for idx, text_encoder in enumerate(text_encoders): + if text_encoder is None: + continue + n = len(token_embeddings[f'txt_encoder_{idx}']) + for i in range(n): + token = trainable_tokens[f'txt_encoder_{idx}'][i] + # Strip any backslashes from the token name: + token = token.replace("/", "_") + embedding = token_embeddings[f'txt_encoder_{idx}'][i] + plot_torch_hist(embedding, global_step, os.path.join(config.output_dir, 'ti_embeddings') , f"enc_{idx}_tokid_{i}: {token}", min_val=-0.05, max_val=0.05, ymax_f = 0.05, color = 'red') - embedding_handler.print_token_info() - if config.is_lora: # plotting this hist for full unet parameters can run OOM - plot_torch_hist(unet_lora_parameters, global_step, config.output_dir, "lora_weights", min_val=-0.4, max_val=0.4, ymax_f = 0.08) - plot_loss(losses, save_path=f'{config.output_dir}/losses.png') - target_std_dict = {f"text_encoder_{idx}_target": embedding_handler.embeddings_settings[f"std_token_embedding_{idx}"].item() for idx in range(len(text_encoders)) if text_encoders[idx] is not None} - plot_token_stds(token_stds, save_path=f'{config.output_dir}/token_stds.png', target_value_dict=target_std_dict) - plot_grad_norms(grad_norms, save_path=f'{config.output_dir}/grad_norms.png') - plot_lrs(optimizer_collection.learning_rate_tracker, save_path=f'{config.output_dir}/learning_rates.png') - plot_curve(prompt_embeds_norms, 'steps', 'norm', 'prompt_embed norms', save_path=f'{config.output_dir}/prompt_embeds_norms.png') - + embedding_handler.print_token_info() + if config.is_lora: # plotting this hist for full unet parameters can run OOM + plot_torch_hist(unet_lora_parameters, global_step, config.output_dir, "lora_weights", min_val=-0.4, max_val=0.4, ymax_f = 0.08) + plot_loss(losses, save_path=f'{config.output_dir}/losses.png') + target_std_dict = {f"text_encoder_{idx}_target": embedding_handler.embeddings_settings[f"std_token_embedding_{idx}"].item() for idx in range(len(text_encoders)) if text_encoders[idx] is not None} + plot_token_stds(token_stds, save_path=f'{config.output_dir}/token_stds.png', target_value_dict=target_std_dict) + plot_grad_norms(grad_norms, save_path=f'{config.output_dir}/grad_norms.png') + plot_lrs(optimizer_collection.learning_rate_tracker, save_path=f'{config.output_dir}/learning_rates.png') + plot_curve(prompt_embeds_norms, 'steps', 'norm', 'prompt_embed norms', save_path=f'{config.output_dir}/prompt_embeds_norms.png') + validation_prompts = render_images( pipe = pipe, render_size = config.validation_img_size, @@ -435,10 +436,10 @@ def train(config: TrainingConfig): images_done += config.train_batch_size global_step += 1 - if global_step % (config.max_train_steps//20) == 0: + if global_step % (config.max_train_steps//50) == 0: progress = (global_step / config.max_train_steps) + 0.05 #print_system_info() - print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r", flush = True) + print(f"\n---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r", flush = True) yield np.min((progress, 1.0)) if global_step > config.max_train_steps: diff --git a/node.py b/node.py index e212e91..e0ea822 100644 --- a/node.py +++ b/node.py @@ -11,6 +11,7 @@ from trainer.config import TrainingConfig, model_paths from trainer.utils.io import clean_filename import folder_paths +import comfy.utils class Eden_LoRa_trainer: @classmethod @@ -37,7 +38,7 @@ class Eden_LoRa_trainer: CATEGORY = "Eden 🌱" RETURN_TYPES = ("IMAGE", "STRING", "STRING", "STRING") - RETURN_NAMES = ("sample_images", "lora_path", "embedding_path", "total_runtime_s") + RETURN_NAMES = ("sample_images", "lora_path", "embedding_path", "final_msg") FUNCTION = "train_lora" def train_lora(self, @@ -87,17 +88,17 @@ class Eden_LoRa_trainer: verbose=True, debug=debug_mode, ) - + + pbar = comfy.utils.ProgressBar(100) + with torch.inference_mode(False): train_generator = train(config=config) while True: try: progress_f = next(train_generator) - print("---- Progress ----") - print(progress_f) + pbar.update_absolute(progress_f * 100) except StopIteration as e: config, output_save_dir = e.value # Capture the return value - print("---- Training finished at StopIteration! ----") break validation_grid_img_path = os.path.join(output_save_dir, "validation_grid.jpg") @@ -107,10 +108,17 @@ class Eden_LoRa_trainer: attributes['job_time_seconds'] = config.job_time print(f"LORA training node finished in {config.job_time:.1f} seconds") - print("---------- Made by Eden.art 🌱 ----------") + print("---------- Made with love by Eden.art 🌱 ----------") + + # safetensors paths: + paths = [os.path.join(output_save_dir, f) for f in os.listdir(output_save_dir) if f.endswith(".safetensors")] - lora_path = os.path.join(output_save_dir, config.name + ".safetensors") - embedding_path = os.path.join(output_save_dir, config.name + "_embeddings.safetensors") + # find the index of the path containing "_embeddings.safetensors": + for i, path in enumerate(paths): + if "_embeddings.safetensors" in path: + embedding_path = path + else: + lora_path = path # Load the grid image: grid_image = Image.open(validation_grid_img_path)