512res

dataset_fromfolder.py

@@ -27,8 +27,8 @@ empty_share = 0.05
 limit = 0
 textemb_full = False
 # Основная процедура обработки
-folder_path = "/workspace/d23"
-save_path = "/workspace/sdxs/datasets/ds23_576"
+folder_path = "/workspace/eshu"
+save_path = "/workspace/sdxs/datasets/eshu_576"
 os.makedirs(save_path, exist_ok=True)
 
 # Функция для очистки CUDA памяти

src/dataset_combine.py

@@ -65,4 +65,4 @@ def combine_datasets(main_dataset_path, datasets_to_add):
     
     return combined
 
-combine_datasets("/workspace/sdxs/datasets/mjnj_576", ["/workspace/sdxs/datasets/ds23_576"])
+combine_datasets("/workspace/sdxs/datasets/576", ["/workspace/sdxs/datasets/eshu_576"])

train.py

@@ -25,12 +25,12 @@ import bitsandbytes as bnb
 import torch.nn.functional as F
 
 # --------------------------- Параметры ---------------------------
-ds_path = "datasets/384"
+ds_path = "datasets/576"
 project = "unet"
-batch_size = 40
-base_learning_rate = 1e-5
-min_learning_rate = 9e-6
-num_epochs = 6
+batch_size = 50
+base_learning_rate = 5e-5
+min_learning_rate = 1e-5
+num_epochs = 20
 # samples/save per epoch
 sample_interval_share = 10
 use_wandb = True
@@ -43,11 +43,14 @@ unet_gradient = True
 clip_sample = False #Scheduler
 fixed_seed = False
 shuffle = True
+dispersive_loss = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.enable_mem_efficient_sdp(True)
 dtype = torch.float32
 save_barrier = 1.03
+dispersive_temperature=0.5
+dispersive_weight=0.25
 percentile_clipping = 97 # Lion
 steps_offset = 1 # Scheduler
 limit = 0
@@ -91,6 +94,93 @@ class AccelerateDispersiveLoss:
         self.activations = []
         self.hooks = []
     
+    def register_hooks(self, model, target_layer="down_blocks.0"):
+        unwrapped_model = self.accelerator.unwrap_model(model)
+        print("=== Поиск слоев в unwrapped модели ===")
+        for name, module in unwrapped_model.named_modules():
+            if target_layer in name:
+                hook = module.register_forward_hook(self.hook_fn)
+                self.hooks.append(hook)
+                print(f"✅ Хук зарегистрирован на: {name}")
+                break
+    
+    def hook_fn(self, module, input, output):
+        
+        if isinstance(output, tuple):
+            activation = output[0]
+        else:
+            activation = output
+            
+        if len(activation.shape) > 2:
+            activation = activation.view(activation.shape[0], -1)
+            
+        self.activations.append(activation.detach())
+        
+    def compute_dispersive_loss(self):
+            if not self.activations:
+                return torch.tensor(0.0, requires_grad=True)
+            
+            local_activations = self.activations[-1].float()
+            
+            batch_size = local_activations.shape[0]
+            if batch_size < 2:
+                return torch.tensor(0.0, requires_grad=True)
+            
+            # Нормализация и вычисление loss
+            sf = local_activations / torch.norm(local_activations, dim=1, keepdim=True)
+            distance = torch.nn.functional.pdist(sf.float(), p=2) ** 2
+            exp_neg_dist = torch.exp(-distance / self.temperature) + 1e-5
+            dispersive_loss = torch.log(torch.mean(exp_neg_dist))
+            
+            # ВАЖНО: он отриц и должен падать
+            return dispersive_loss
+    
+    def compute_dispersive_loss2(self):
+        # Если нет активаций, возвращаем 0
+        if not self.activations:
+            return torch.tensor(0.0, device=self.accelerator.device, requires_grad=True)
+        
+        # Работаем только с локальными активациями главного процесса
+        activations = self.activations[-1].float()
+        
+        batch_size = activations.shape[0]
+        if batch_size < 2:
+            return torch.tensor(0.0, device=self.accelerator.device, requires_grad=True)
+        
+        # Нормализация
+        norm = torch.norm(activations, dim=1, keepdim=True).clamp(min=1e-12)
+        sf = activations / norm
+        
+        # Вычисляем расстояния
+        distance = torch.nn.functional.pdist(sf, p=2)
+        distance = distance.clamp(min=1e-12)
+        distance_squared = distance ** 2
+        
+        # Вычисляем loss с клиппингом для стабильности
+        exp_neg_dist = torch.exp((-distance_squared / self.temperature).clamp(min=-20, max=20))
+        exp_neg_dist = exp_neg_dist + 1e-12
+        
+        mean_exp = torch.mean(exp_neg_dist)
+        dispersive_loss = torch.log(mean_exp.clamp(min=1e-12))
+        
+        return dispersive_loss
+    
+    def clear_activations(self):
+        self.activations.clear()
+    
+    def remove_hooks(self):
+        for hook in self.hooks:
+            hook.remove()
+        self.hooks.clear()
+            
+class AccelerateDispersiveLoss2:
+    def __init__(self, accelerator, temperature=0.5, weight=0.5):
+        self.accelerator = accelerator
+        self.temperature = temperature
+        self.weight = weight
+        self.activations = []
+        self.hooks = []
+    
     def register_hooks(self, model, target_layer="down_blocks.0"):
         # Получаем "чистую" модель без DDP wrapper'а
         unwrapped_model = self.accelerator.unwrap_model(model)
@@ -455,12 +545,13 @@ if os.path.isdir(latest_checkpoint):
     if hasattr(torch.nn.functional, "get_flash_attention_available"):
          print(f"torch.nn.functional.get_flash_attention_available(): {torch.nn.functional.get_flash_attention_available()}")
    
-    # Регистрируем хуки на модел
-    dispersive_hook = AccelerateDispersiveLoss(
-        accelerator=accelerator,
-        temperature=2,
-        weight=0.25
-    )
+    # Регистрируем хук на модел
+    if dispersive_loss:
+        dispersive_hook = AccelerateDispersiveLoss(
+            accelerator=accelerator,
+            temperature=dispersive_temperature,
+            weight=dispersive_weight
+        )
     
 if torch_compile:
     print("compiling")
@@ -590,8 +681,9 @@ else:
     lr_scheduler = LambdaLR(optimizer, lambda step: lr_schedule(step) / base_learning_rate)
     unet, optimizer, lr_scheduler = accelerator.prepare(unet, optimizer, lr_scheduler)
 
- # Регистрация хуков ПОСЛЕ prepare
-dispersive_hook.register_hooks(unet, "down_blocks.2")
+# Регистрация хуков ПОСЛЕ prepare
+if dispersive_loss:
+    dispersive_hook.register_hooks(unet, "down_blocks.2")
 
 # --------------------------- Фиксированные семплы для генерации ---------------------------
 # Примеры фиксированных семплов по размерам
@@ -611,9 +703,7 @@ def generate_and_save_samples(fixed_samples_cpu, step):
     original_model = None # Инициализируем, чтобы finally не ругался
     try:
 
-        original_model = accelerator.unwrap_model(unet)
-        original_model = original_model.to(dtype = dtype)
-        original_model.eval() 
+        original_model = accelerator.unwrap_model(unet).eval()
 
         vae.to(device=device, dtype=dtype)
         vae.eval()
@@ -705,9 +795,6 @@ def generate_and_save_samples(fixed_samples_cpu, step):
     
     finally:        
         vae.to("cpu") # Перемещаем VAE обратно на CPU
-        original_model = original_model.to(dtype = dtype)
-        if original_model is not None:
-            del original_model
         # Очистка переменных, которые являются тензорами и были созданы в функции
         for var in list(locals().keys()):
             if isinstance(locals()[var], torch.Tensor):
@@ -721,6 +808,7 @@ if accelerator.is_main_process:
     if save_model:
         print("Генерация сэмплов до старта обучения...")
         generate_and_save_samples(fixed_samples,0)
+accelerator.wait_for_everyone()
 
 # Модифицируем функцию сохранения модели для поддержки LoRA
 def save_checkpoint(unet,variant=""):
@@ -775,7 +863,8 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             noisy_latents = scheduler.add_noise(latents, noise, timesteps)
 
             # Очищаем активации перед forward pass
-            dispersive_hook.clear_activations()
+            if dispersive_loss:
+                dispersive_hook.clear_activations()
         
             # Используем целевое значение
             model_pred = unet(noisy_latents, timesteps, embeddings).sample
@@ -787,8 +876,12 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             # Dispersive Loss
             #Идентичные векторы: Loss = -0.0000
             #Ортогональные векторы: Loss = -3.9995
-            with torch.cuda.amp.autocast(enabled=False):
-                dispersive_loss = dispersive_hook.weight * dispersive_hook.compute_dispersive_loss()
+            if dispersive_loss:
+                with torch.cuda.amp.autocast(enabled=False):
+                    dispersive_loss = dispersive_hook.weight * dispersive_hook.compute_dispersive_loss()
+                    if torch.isnan(dispersive_loss) or torch.isinf(dispersive_loss):
+                        print(f"Rank {accelerator.process_index}: Found nan/inf in dispersive_loss: {total_loss}")
+                        break
 
             # Итоговый loss
             # dispersive_loss должен падать и тотал падать - поэтому плюс
@@ -800,17 +893,20 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
                 save_model = False
                 break
             
-            # Делаем backward через Accelerator
+            if (global_step % 100 == 0) or (global_step % sample_interval == 0):
+                accelerator.wait_for_everyone()
+                
+            # Делаем backward через Accelerator    
             accelerator.backward(total_loss)
 
             if (global_step % 100 == 0) or (global_step % sample_interval == 0):
                 accelerator.wait_for_everyone()
-    
+                
             grad = 0.0
             if not fbp:
                 if accelerator.sync_gradients:
-                    with torch.cuda.amp.autocast(enabled=False):
-                        grad = accelerator.clip_grad_norm_(unet.parameters(), 1.)
+                    #with torch.cuda.amp.autocast(enabled=False):
+                    grad = accelerator.clip_grad_norm_(unet.parameters(), 1.)
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad(set_to_none=True)
@@ -873,7 +969,8 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             wandb.log({"epoch_loss": avg_epoch_loss, "epoch": epoch+1})
 
 # Завершение обучения - сохраняем финальную модель
-dispersive_hook.remove_hooks()
+if dispersive_loss:
+    dispersive_hook.remove_hooks()
 if accelerator.is_main_process:
     print("Обучение завершено! Сохраняем финальную модель...")
     # Сохраняем основную модель

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7905a5f9fc1ead936613fd66ef66675ab7cde73128c49208ee1c02668dcb8527
+oid sha256:37bee8e3947ce359ec56fedc8c30322465ad3a69e62d4a3964b9c433d975c34c
 size 7014306128