changes.patch

diff --git a/fine_tune.py b/fine_tune_new.py
index 2658d55..0fb94a0 100755
--- a/fine_tune.py
+++ b/fine_tune_new.py
@@ -151,8 +151,38 @@ class FineTuningDataset(torch.utils.data.Dataset):
 
       if self.shuffle_caption:
         tokens = caption.strip().split(",")
-        random.shuffle(tokens)
-        caption = ",".join(tokens).strip()
+        copy = tokens[1:]
+        random.shuffle(copy)
+        
+        while ("" in copy): ##removing empty tags
+          copy.remove("")
+        
+        temphold = [] ##Temporarily storing all xgirl/yboys tags and removing them from copy list
+        for i in copy:
+          if i.strip()[1:] in ["boy", "boys", "girl", "girls"]:
+            temphold.append(i)
+            copy.remove(i)
+        
+
+        
+        value = (len(copy))
+        
+        if value > 4:
+          deltags = random.randint(1, value-2)
+          copy = copy[deltags:]
+        else: deltags = 0
+        
+        for i in temphold:
+          copy.append(i)
+        random.shuffle(copy)
+        
+
+        tokens[1:] = copy
+        #print(tokens)
+
+        
+
+      caption = ",".join(tokens).strip()
 
       captions.append(caption)
 
@@ -341,7 +371,7 @@ def train(args):
 
   # lr schedulerを用意する
   lr_scheduler = diffusers.optimization.get_scheduler(
-      "constant", optimizer, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
+      "constant", optimizer=optimizer, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
 
   # acceleratorがなんかよろしくやってくれるらしい
   if fine_tuning:
@@ -353,21 +383,22 @@ def train(args):
   else:
     unet, hypernetwork, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
         unet, hypernetwork, optimizer, train_dataloader, lr_scheduler)
-
+  
+  
   # resumeする
   if args.resume is not None:
     print(f"resume training from state: {args.resume}")
     accelerator.load_state(args.resume)
 
   # epoch数を計算する
-  num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+  num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) 
   num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
   # 学習する
   total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
   print("running training / 学習開始")
   print(f"  num examples / サンプル数: {train_dataset.images_count}")
-  print(f"  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
+  print(f"  num batches per epoch / 1epochのバッチ数: {num_update_steps_per_epoch}")
   print(f"  num epochs / epoch数: {num_train_epochs}")
   print(f"  batch size per device / バッチサイズ: {args.train_batch_size}")
   print(f"  total train batch size (with parallel & distributed) / 総バッチサイズ（並列学習含む）: {total_batch_size}")
@@ -377,7 +408,7 @@ def train(args):
   progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process, desc="steps")
   global_step = 0
 
-  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
+  noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False)
 
   if accelerator.is_main_process:
     accelerator.init_trackers("finetuning" if fine_tuning else "hypernetwork")