add conv2 to lora layers

scripts/create_hyperparam_sweep.py

@@ -35,7 +35,7 @@ def hamming_distance(dict1, dict2):
 #######################################################################################
 
 # Setup the base experiment config:
-exp_name             = "unet_adam_object"
+exp_name             = "sd15_face_sweep"
 caption_prefix       = ""
 mask_target_prompts  = ""
 n_exp                = 200  # how many random experiment settings to generate
@@ -55,14 +55,14 @@ hyperparameters = {
     ],
     "concept_mode": ['face'],
     "seed": [0],
-    "resolution": [512,640,768,1024],
-    "train_batch_size": [3],
+    "resolution": [512,640,768],
+    "train_batch_size": [4],
     "n_sample_imgs": [6],
     "max_train_steps": [400,800],
     "checkpointing_steps": [100],
-    "gradient_accumulation_steps": [1],    
+    "gradient_accumulation_steps": [1],
 
-    "n_tokens": [2],
+    "n_tokens": [1, 2],
     "ti_lr": [0.001],
     "ti_weight_decay": [0.0005],
     "l1_penalty": [0.0],
@@ -72,8 +72,7 @@ hyperparameters = {
     "tok_cond_reg_w": [0.01e-5, 2.5e-5],
 
     "unet_prodigy_growth_factor": [1.05],
-    "unet_lr_warmup_steps": [100,200,400],
-    "unet_lr": [1.0e-4, 3e-4, 1e-3, 3e-3],
+    "unet_lr": [1.0e-4, 3e-4, 1e-3],
     "prodigy_d_coef": [1.0],
     "lora_weight_decay": [0.001],
     "lora_rank": [6, 12, 24, 48],
@@ -146,7 +145,7 @@ 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 -c {os.path.join(folder_path, json_file)}\n"
+            command = f"python main.py {os.path.join(folder_path, json_file)}\n"
             sh_file.write(command)
 
 generate_sh_script(config_output_dir, output_sh_path)

trainer/config.py

@@ -102,7 +102,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.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}"
         os.makedirs(self.output_dir, exist_ok=True)
 
         if self.seed is None:

trainer/preprocess.py

@@ -713,7 +713,7 @@ def load_and_save_masks_and_captions(
 
     if config.validation_img_size is None:
         config.validation_img_size = [0, 0]
-        multiplier = 2.0 if config.sd_model_version == "sdxl" else 1.25
+        multiplier = 2.0 if config.sd_model_version == "sdxl" else 1.0
         config.validation_img_size[0] = config.train_img_size[0] * multiplier
         config.validation_img_size[1] = config.train_img_size[1] * multiplier
     elif isinstance(config.validation_img_size, int):

trainer/utils/utils.py

@@ -90,25 +90,28 @@ def pick_best_gpu_id():
 
 
 def plot_torch_hist(parameters, step, checkpoint_dir, name, bins=100, min_val=-1, max_val=1, ymax_f = 0.75, color = 'blue'):
-    os.makedirs(checkpoint_dir, exist_ok=True)
+    try:
+        os.makedirs(checkpoint_dir, exist_ok=True)
 
-    # Flatten and concatenate all parameters into a single tensor
-    all_params = torch.cat([p.data.view(-1) for p in parameters])
-    norm = torch.norm(all_params)
+        # Flatten and concatenate all parameters into a single tensor
+        all_params = torch.cat([p.data.view(-1) for p in parameters])
+        norm = torch.norm(all_params)
 
-    # Convert to CPU for plotting
-    all_params_cpu = all_params.cpu().float().numpy()
+        # Convert to CPU for plotting
+        all_params_cpu = all_params.cpu().float().numpy()
 
-    # Plot histogram
-    plt.figure()
-    plt.hist(all_params_cpu, bins=bins, density=False, color = color)
-    plt.ylim(0, ymax_f * len(all_params_cpu.flatten()))
-    plt.xlim(min_val, max_val)
-    plt.xlabel('Weight Value')
-    plt.ylabel('Count')
-    plt.title(f'{name} (std: {np.std(all_params_cpu):.5f}, norm: {norm:.3f}, step {step:03d})')
-    plt.savefig(f"{checkpoint_dir}/{name}_hist_{step:04d}.png")
-    plt.close()
+        # Plot histogram
+        plt.figure()
+        plt.hist(all_params_cpu, bins=bins, density=False, color = color)
+        plt.ylim(0, ymax_f * len(all_params_cpu.flatten()))
+        plt.xlim(min_val, max_val)
+        plt.xlabel('Weight Value')
+        plt.ylabel('Count')
+        plt.title(f'{name} (std: {np.std(all_params_cpu):.5f}, norm: {norm:.3f}, step {step:03d})')
+        plt.savefig(f"{checkpoint_dir}/{name}_hist_{step:04d}.png")
+        plt.close()
+    except:
+        print(f'Error plotting {name} histogram')
 
 def plot_curve(value_dict, xlabel, ylabel, title, save_path, log_scale = False, y_lims = None):
     plt.figure()