From 87d656c9e4f78a563ebc0e8dde0dbefb5c024ac3 Mon Sep 17 00:00:00 2001 From: aiXander Date: Mon, 11 Mar 2024 22:18:56 -0700 Subject: [PATCH] more refactoring --- io_utils.py | 2 +- predict.py | 1 + trainer_pti.py | 6 +++--- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/io_utils.py b/io_utils.py index 44abf70..9e034aa 100755 --- a/io_utils.py +++ b/io_utils.py @@ -25,7 +25,7 @@ MODEL_INFO = { def download_weights(url, dest): start = time.time() print("downloading url: ", url) - print("downloading to: ", dest) + print("downloading to: ", dest, '...') # Make sure the destination directory exists dest_dir = os.path.dirname(dest) diff --git a/predict.py b/predict.py index 3e8080e..483e156 100755 --- a/predict.py +++ b/predict.py @@ -374,6 +374,7 @@ class Predictor(BasePredictor): scale_lr=False, allow_tf32=True, mixed_precision="bf16", + #mixed_precision="fp16", # this 100% breaks training... Figure out why!! device="cuda:0", lora_rank=lora_rank, is_lora=is_lora, diff --git a/trainer_pti.py b/trainer_pti.py index a83c7e1..2a79a93 100755 --- a/trainer_pti.py +++ b/trainer_pti.py @@ -318,7 +318,7 @@ def render_images(pipeline, lora_path, train_step, seed, is_lora, pretrained_mod text_encoder_one, text_encoder_two, vae, - unet) = load_models(pretrained_model, device, torch.float16) + unet) = load_models(pretrained_model, device, torch.float16) #, keep_vae_float32 = True) pipeline = pipeline.to(device) pipeline = patch_pipe_with_lora(pipeline, lora_path) @@ -464,7 +464,7 @@ def main( text_encoder_two, vae, unet, - ) = load_models(pretrained_model, device, weight_dtype, keep_vae_float32 = True) + ) = load_models(pretrained_model, device, weight_dtype)#, keep_vae_float32 = True) # Initialize new tokens for training. embedding_handler = TokenEmbeddingsHandler( @@ -854,10 +854,10 @@ def main( if (global_step % checkpointing_steps == 0): output_save_dir = f"{checkpoint_dir}/checkpoint-{global_step}" save(output_save_dir, global_step, unet, embedding_handler, token_dict, args_dict, seed, is_lora, unet_lora_parameters, unet_param_to_optimize_names) - validation_prompts = render_images(pipe, output_save_dir, global_step, seed, is_lora, pretrained_model, n_imgs = 4, debug=debug) last_save_step = global_step if debug: + validation_prompts = render_images(pipe, output_save_dir, global_step, seed, is_lora, pretrained_model, n_imgs = 4, debug=debug) token_embeddings = embedding_handler.get_trainable_embeddings() for i, token_embeddings_i in enumerate(token_embeddings): plot_torch_hist(token_embeddings_i[0], global_step, output_dir, f"embeddings_weights_token_0_{i}", min_val=-0.05, max_val=0.05, ymax_f = 0.05)