From 4673ec5ef4b488aa88444f67b87bc7461ff8e0db Mon Sep 17 00:00:00 2001 From: mayukhdeb Date: Sat, 16 Mar 2024 06:47:55 -0700 Subject: [PATCH] nuke args_dict (everything working) --- predict.py | 27 ++++----------------------- 1 file changed, 4 insertions(+), 23 deletions(-) diff --git a/predict.py b/predict.py index ba70d82..bf1fb5b 100755 --- a/predict.py +++ b/predict.py @@ -403,30 +403,11 @@ class Predictor(BasePredictor): output_save_dir, validation_prompts = e.value # Capture the return value break - if not debug: - keys_to_keep = [ - "name", - "checkpoint", - "concept_mode", - "input_images", - "num_training_images", - "seed", - "resolution", - "max_train_steps", - "lora_rank", - "trigger_text", - "left_right_flip_augmentation", - "run_name", - "trainig_captions"] - args_dict = {k: v for k, v in args_dict.items() if k in keys_to_keep} - - args_dict["grid_prompts"] = validation_prompts - # save final training_args: final_args_dict_path = os.path.join(output_dir, "training_args.json") - with open(final_args_dict_path, "w") as f: - json.dump(args_dict, f, indent=4) - + config.save_as_json( + final_args_dict_path + ) validation_grid_img_path = os.path.join(output_save_dir, "validation_grid.jpg") out_path = f"{clean_filename(name)}_eden_concept_lora_{int(time.time())}.tar" directory = cogPath(output_save_dir) @@ -454,4 +435,4 @@ class Predictor(BasePredictor): else: # clear the output_directory to avoid running out of space on the machine: #shutil.rmtree(output_dir) - yield CogOutput(files=[cogPath(out_path)], name=name, thumbnails=[cogPath(validation_grid_img_path)], attributes=args_dict, isFinal=True, progress=1.0) \ No newline at end of file + yield CogOutput(files=[cogPath(out_path)], name=name, thumbnails=[cogPath(validation_grid_img_path)], attributes=config.dict(), isFinal=True, progress=1.0) \ No newline at end of file