diff --git a/main.py b/main.py index 1c65582..523e8db 100755 --- a/main.py +++ b/main.py @@ -551,7 +551,10 @@ def main( with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf: zipdir(parent_dir, zipf) - return output_save_dir, validation_prompts + config.job_time = time.time() - config.start_time + config.training_attributes["validation_prompts"] = validation_prompts + + return config, output_save_dir diff --git a/predict.py b/predict.py index 01b10e8..6c0aa60 100755 --- a/predict.py +++ b/predict.py @@ -140,39 +140,34 @@ class Predictor(BasePredictor): bs=6: 8.0 imgs/s, """ - start_time = time.time() - - config = TrainingConfig( - output_dir="path/to/output", - name="my_training", - lora_training_urls="https://example.com/lora", - concept_mode="face", - sd_model_version="sdxl", - # Add other parameters as needed - ) - - out_root_dir = "lora_models" - - if seed is None: - seed = np.random.randint(0, 2**32 - 1) - - # Try to make the training reproducible: - seed_everything(seed = seed) - - print(f"cog:predict:train_lora:{concept_mode}") - + print("cog:predict starting new training job...") if not debug: yield CogOutput(name=name, progress=0.0) - - # Initialize pretrained_model dictionary - pretrained_model = pretrained_models[sd_model_version] - - # hardcoded for now: - token_list = [f"TOK:{n_tokens}"] - - inserting_list_tokens, token_dict = obtain_inserting_list_tokens(token_list=token_list) - - + + config = TrainingConfig( + name=name, + lora_training_urls=lora_training_urls, + concept_mode=concept_mode, + sd_model_version=sd_model_version, + seed=seed, + resolution=resolution, + train_batch_size=train_batch_size, + max_train_steps=max_train_steps, + checkpointing_steps=checkpointing_steps, + is_lora=is_lora, + prodigy_d_coef=prodigy_d_coef, + ti_lr=ti_lr, + ti_weight_decay=ti_weight_decay, + lora_weight_decay=lora_weight_decay, + l1_penalty=l1_penalty, + lora_rank=lora_rank, + caption_model=caption_model, + n_tokens=n_tokens, + verbose=verbose, + debug=debug, + off_ratio_power=off_ratio_power + ) + if XANDER_EXPERIMENT: # overwrite some settings for experimentation: lora_param_scaler = 0.1 @@ -239,73 +234,9 @@ class Predictor(BasePredictor): merge_datasets(input_dir1, input_dir2, input_dir, token_dict.keys()) else: # normal, single token run: - - output_dir = os.path.join(out_root_dir, run_name) - - if not debug: - yield CogOutput(name=name, progress=0.05) - - config = TrainingConfig( - name=name, - pretrained_model=pretrained_model, - lora_training_urls=lora_training_urls, - concept_mode=concept_mode, - sd_model_version=sd_model_version, - seed=seed, - resolution=resolution, - train_batch_size=train_batch_size, - num_train_epochs=num_train_epochs, - max_train_steps=max_train_steps, - checkpointing_steps=checkpointing_steps, - gradient_accumulation_steps=gradient_accumulation_steps, - is_lora=is_lora, - prodigy_d_coef=prodigy_d_coef, - ti_lr=ti_lr, - ti_weight_decay=ti_weight_decay, - lora_weight_decay=lora_weight_decay, - l1_penalty=l1_penalty, - lora_param_scaler=lora_param_scaler, - snr_gamma=snr_gamma, - lora_rank=lora_rank, - caption_prefix=caption_prefix, - caption_model=caption_model, - left_right_flip_augmentation=left_right_flip_augmentation, - augment_imgs_up_to_n=augment_imgs_up_to_n, - n_tokens=n_tokens, - mask_target_prompts=mask_target_prompts, - crop_based_on_salience=crop_based_on_salience, - use_face_detection_instead=use_face_detection_instead, - clipseg_temperature=clipseg_temperature, - verbose=verbose, - run_name=run_name, - debug=debug, - hard_pivot=hard_pivot, - off_ratio_power=off_ratio_power, - allow_tf32 = True, - weight_type="bf16", - inserting_list_tokens=inserting_list_tokens, - token_dict=token_dict, - device="cuda:0", - output_dir=output_dir, - scale_lr=False, - crops_coords_top_left_h = 0, - crops_coords_top_left_w = 0, - do_cache = True, - unet_learning_rate = 1.0, - lr_scheduler = "constant", - lr_warmup_steps = 50, - lr_num_cycles = 1, - lr_power = 1.0, - dataloader_num_workers = 0, - ) - - config.save_as_json( - os.path.join(output_dir, "training_args.json") - ) - - train_generator = main( - config=config - ) + pass + + train_generator = main(config=config) while True: try: @@ -313,7 +244,7 @@ class Predictor(BasePredictor): if not debug: yield CogOutput(name=name, progress=np.round(progress_f, 2)) except StopIteration as e: - output_save_dir, validation_prompts = e.value # Capture the return value + config, output_save_dir = e.value # Capture the return value break # save final training_args: @@ -336,9 +267,8 @@ class Predictor(BasePredictor): tar.add("instructions_README.md", arcname="README.md") attributes = {} - attributes['grid_prompts'] = validation_prompts - runtime = time.time() - start_time - attributes['job_time_seconds'] = runtime + attributes['grid_prompts'] = config.training_attributes["validation_prompts"] + attributes['job_time_seconds'] = config.job_time print(f"LORA training finished in {runtime:.1f} seconds") print(f"Returning {out_path}") diff --git a/preprocess.py b/preprocess.py index 329689e..d44df1e 100755 --- a/preprocess.py +++ b/preprocess.py @@ -255,60 +255,46 @@ def cleanup_prompts_with_chatgpt( seed, # seed for chatgpt reproducibility verbose = True): - if concept_mode == "object_injection": - chat_gpt_prompt_1 = textwrap.dedent(""" - I have a set of images, each containing the same concept / figure. I have the following (poor) descriptions for each image: - """) - - chat_gpt_prompt_2 = textwrap.dedent(""" - I want you to: - 1. Find a good, short name/description of the single central concept that's in all the images. This [Concept Name] might eg already be present in the descriptions above, pick the most obvious name or words that would fit in all descriptions. - 2. Insert the text "TOK, [Concept Name]" into all the descriptions above by rephrasing them where needed to naturally contain the text TOK, [Concept Name] while keeping as much of the description as possible. - - Reply by first stating the "Concept Name:", followed by an enumerated list (using "-") of all the revised "Descriptions:". - """) - if concept_mode == "object": chat_gpt_prompt_1 = textwrap.dedent(""" - Analyze a set of (poor) image descriptions each featuring the same concept, figure or thing. - Tasks: - 1. Deduce a concise, fitting name for the concept that is visually descriptive (Concept Name). - 2. Substitute the concept in each description with the placeholder "TOK", rearranging or adjusting the text where needed. Hallucinate TOK into the description if necessary (but dont mention when doing so, simply provide the final description)! - 3. Streamline each description to its core elements, ensuring clarity and mandatory inclusion of the placeholder string "TOK". - The descriptions are: - """) + Analyze a set of (poor) image descriptions each featuring the same concept, figure or thing. + Tasks: + 1. Deduce a concise, fitting name for the concept that is visually descriptive (Concept Name). + 2. Substitute the concept in each description with the placeholder "TOK", rearranging or adjusting the text where needed. Hallucinate TOK into the description if necessary (but dont mention when doing so, simply provide the final description)! + 3. Streamline each description to its core elements, ensuring clarity and mandatory inclusion of the placeholder string "TOK". + The descriptions are: + """) chat_gpt_prompt_2 = textwrap.dedent(""" - Respond with the chosen "Concept Name:" followed by a list (using "-") of all the revised descriptions, each mentioning "TOK". - """) + Respond with the chosen "Concept Name:" followed by a list (using "-") of all the revised descriptions, each mentioning "TOK". + """) elif concept_mode == "face": chat_gpt_prompt_1 = textwrap.dedent(""" - Analyze a set of (poor) image descriptions, each featuring a person named TOK. - Tasks: - 1. Rewrite each description, ensuring it refers only to a single person or character. - 2. Integrate "a photo of TOK" naturally into each description, rearranging or adjusting where needed. - 3. Streamline each description to its core elements, ensuring clarity and mandatory inclusion of "TOK". - The descriptions are: - """) + Analyze a set of (poor) image descriptions, each featuring a person named TOK. + Tasks: + 1. Rewrite each description, ensuring it naturally refers to "TOK", rearranging or adjusting where needed. + 2. Streamline each description to its core elements, ensuring clarity and mandatory inclusion of "TOK". + The descriptions are: + """) chat_gpt_prompt_2 = textwrap.dedent(""" - Respond with "Concept Name: TOK" followed by a list (using "-") of all the revised descriptions, each mentioning "a photo of TOK". - """) + Respond with "Concept Name: TOK" followed by a list (using "-") of all the revised descriptions, each mentioning "a photo of TOK". + """) elif concept_mode == "style": chat_gpt_prompt_1 = textwrap.dedent(""" - Analyze a set of (poor) image descriptions, each featuring the same style named TOK. - Tasks: - 1. Rewrite each description to focus solely on the TOK style. - 2. Integrate "in the style of TOK" naturally into each description, typically at the beginning. - 3. Summarize each description to its core elements, ensuring clarity and mandatory inclusion of "TOK". - The descriptions are: - """) + Analyze a set of (poor) image descriptions, each featuring an example a common aesthetic style named TOK. + Tasks: + 1. Rewrite each description to focus solely on the non-stylistic contents of the images like characters, objects, colors, scene, context etc. + 2. Integrate "in the style of TOK" naturally into each description, typically at the beginning. + 3. Summarize each description to its core elements, ensuring clarity and mandatory inclusion of "TOK". + The descriptions are: + """) chat_gpt_prompt_2 = textwrap.dedent(""" - Respond with "Style Name: TOK" followed by a list (using "-") of all the revised descriptions, each mentioning "in the style of TOK". - """) + Respond with "Style Name: TOK" followed by a list (using "-") of all the revised descriptions, each mentioning "in the style of TOK". + """) final_chatgpt_prompt = chat_gpt_prompt_1 + "\n- " + "\n- ".join(prompts) + "\n\n" + chat_gpt_prompt_2 print("Final chatgpt prompt:") diff --git a/trainer/config.py b/trainer/config.py index 6fcc0e5..d9b489b 100644 --- a/trainer/config.py +++ b/trainer/config.py @@ -69,6 +69,8 @@ class TrainingConfig(BaseModel): dataloader_num_workers: int = 0 training_attributes: dict = {} aspect_ratio_bucketing: bool = False + start_time: float = 0.0 + job_time: float = 0.0 def __init__(self, **data): super().__init__(**data) @@ -100,6 +102,7 @@ class TrainingConfig(BaseModel): gpu_id = pick_best_gpu_id() self.device = f'cuda:{gpu_id}' + self.start_time = time.time() @classmethod def from_json(cls, file_path: str): diff --git a/training_args_style.json b/training_args_style.json index ba84a48..0baf648 100644 --- a/training_args_style.json +++ b/training_args_style.json @@ -8,9 +8,9 @@ "resolution": 512, "validation_img_size": [1024, 1024], "train_batch_size": 4, - "n_sample_imgs": 6, - "max_train_steps": 600, - "checkpointing_steps": 200, + "n_sample_imgs": 4, + "max_train_steps": 60, + "checkpointing_steps": 30, "gradient_accumulation_steps": 1, "n_tokens": 3, "ti_lr": 0.001,