cleanup #3

2024-04-08 00:02:08 -07:00
parent 868e0e1b2f
commit ea470f5de7
5 changed files with 68 additions and 146 deletions
@@ -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



@@ -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}")
@@ -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:")
@@ -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):
@@ -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,