2512

dataset.py

@@ -27,8 +27,8 @@ step = 32 #64
 empty_share = 0.0
 limit = 0
 # Основная процедура обработки
-folder_path = "/workspace/ds" #alchemist"
-save_path = "/workspace/768" #"alchemist"
+folder_path = "/workspace/alchemist" #alchemist"
+save_path = "/workspace/sdxs/datasets/alchemist" #"alchemist"
 os.makedirs(save_path, exist_ok=True)
 
 # Функция для очистки CUDA памяти
@@ -394,6 +394,13 @@ print(f"Всего найдено {len(image_paths)} изображений")
 # Обработка с чанкованием
 process_in_chunks(image_paths, text_paths, width, height, chunk_size=20000, batch_size=batch_size)
 
+# Удаление  папки
+try:
+    shutil.rmtree(folder_path)
+    print(f"✅ Папка {folder_path} успешно удалена")
+except Exception as e:
+    print(f"⚠️ Ошибка при удалении папки: {e}")
+    
 # Объединение чанков в финальный датасет
 combine_chunks(temp_path, save_path)
 

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a85017c24383c5064153544e98810e64a43a2d9c1b0ce79362f101391f84968b
-size 1723522
+oid sha256:fb11577f5bc95ddaec1ac5dfbd4090581c6ea31c5c5e0bbc1d616d5948ec98aa
+size 9700840

train.py

@@ -1,4 +1,4 @@
-from comet_ml import Experiment
+#from comet_ml import Experiment
 import os
 import math
 import torch
@@ -29,14 +29,14 @@ from transformers import AutoTokenizer, AutoModel
 # --------------------------- Параметры ---------------------------
 ds_path = "/workspace/sdxs/datasets/768"
 project = "unet"
-batch_size = 256
-base_learning_rate = 3e-5
-min_learning_rate = 2.5e-5
-num_epochs = 10
-sample_interval_share = 20
+batch_size = 60
+base_learning_rate = 2.7e-5
+min_learning_rate = 2.7e-5
+num_epochs = 100
+sample_interval_share = 5
 max_length = 192
-use_wandb = False
-use_comet_ml = True
+use_wandb = True
+use_comet_ml = False
 save_model = True
 use_decay = True
 fbp = False
@@ -51,11 +51,11 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 #torch.backends.cuda.enable_mem_efficient_sdp(False)
 dtype = torch.float32
-save_barrier = 1.005
+save_barrier = 1.006
 warmup_percent = 0.01
 percentile_clipping = 98
 betta2 = 0.998
-eps = 1e-6
+eps = 1e-7
 clip_grad_norm = 1.0
 limit = 0
 checkpoints_folder = ""
@@ -171,7 +171,7 @@ def encode_texts(texts, max_length=max_length):
         batch_size = hidden.shape[0]
         pooled = hidden[torch.arange(batch_size, device=hidden.device), sequence_lengths]
 
-        #return hidden, attention_mask, pooled
+        #return hidden, attention_mask
         # --- НОВАЯ ЛОГИКА: ОБЪЕДИНЕНИЕ ДЛЯ КРОСС-ВНИМАНИЯ ---
         # 1. Расширяем пулинг-вектор до последовательности [B, 1, 1024]
         pooled_expanded = pooled.unsqueeze(1) 
@@ -186,7 +186,7 @@ def encode_texts(texts, max_length=max_length):
         # torch.ones((batch_size, 1), device=device) создает маску [B, 1] со значениями 1.
         new_attention_mask = torch.cat([torch.ones((batch_size, 1), device=device), attention_mask], dim=1)
         
-        return new_encoder_hidden_states, new_attention_mask, pooled
+        return new_encoder_hidden_states, new_attention_mask
 
 shift_factor = getattr(vae.config, "shift_factor", 0.0)
 if shift_factor is None: shift_factor = 0.0
@@ -287,9 +287,9 @@ def get_fixed_samples_by_resolution(dataset, samples_per_group=1):
         texts = [item["text"] for item in samples_data]
         
         # Кодируем тексты на лету, чтобы получить маски и пулинг
-        embeddings, masks, pools = encode_texts(texts)
+        embeddings, masks = encode_texts(texts)
         
-        fixed_samples[size] = (latents, embeddings, masks, pools, texts)
+        fixed_samples[size] = (latents, embeddings, masks, texts)
     
     print(f"Создано {len(fixed_samples)} групп фиксированных семплов по разрешениям")
     return fixed_samples
@@ -309,19 +309,19 @@ def collate_fn_simple(batch):
     texts = [
         "" if t.lower().startswith("zero") 
         else "" if random.random() < 0.05
-        else t[1:].lstrip() if t.startswith(".") 
-        else t 
+        else t[1:].lstrip() if t.startswith(".")
+        else t.replace("The image shows ", "").replace("The image is ", "").replace("This image captures ","").strip()
         for t in raw_texts
     ]
     
     # 3. Кодируем на лету
-    # Возвращает: hidden (B, L, D), mask (B, L), pooled (B, D)
-    embeddings, attention_mask, pooled_embeddings = encode_texts(texts)
+    # Возвращает: hidden (B, L, D), mask (B, L)
+    embeddings, attention_mask = encode_texts(texts)
     
     # attention_mask от токенизатора уже имеет нужный формат, но на всякий случай приведем к long
     attention_mask = attention_mask.to(dtype=torch.int64)
 
-    return latents, embeddings, attention_mask, pooled_embeddings
+    return latents, embeddings, attention_mask
 
 batch_sampler = DistributedResolutionBatchSampler(
         dataset=dataset,
@@ -375,8 +375,70 @@ def create_optimizer(name, params):
         )
     elif name == "adam":
         return torch.optim.AdamW(
-            params, lr=base_learning_rate, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01
+            params, lr=base_learning_rate, betas=(0.9, betta2), eps=1e-8, weight_decay=0.01
         )
+    elif name == "adamuon":
+        from pytorch_optimizer import AdaMuon
+        # 1. РАЗДЕЛЕНИЕ ПАРАМЕТРОВ
+        # Разделяем параметры на те, которые Muon-совместимы (ndim >= 2), 
+        # и те, которые лучше оставить для AdamW8bit (ndim < 2).
+        
+        muon_params = []
+        adam8bit_params = []
+
+        if isinstance(params, (list, tuple)):
+            all_params = params
+        else:
+            # Если params - это генератор (unet.parameters()), преобразуем его в list
+            all_params = list(params)
+        
+        for p in all_params:
+            if p.requires_grad:
+                if p.ndim >= 2:
+                    muon_params.append(p)
+                else:
+                    adam8bit_params.append(p)
+        
+        # 2. СОЗДАНИЕ ГРУПП ПАРАМЕТРОВ
+        param_groups = []
+        
+        # --- Группа 1: AdaMuon (для матриц/сверток) ---
+        if muon_params:
+            print(f"Используем AdaMuon для {len(muon_params)} параметров (ndim >= 2)")
+            param_groups.append({
+                'params': muon_params,
+                'optimizer': 'AdaMuon',              # Явно указываем оптимизатор
+                'use_muon': True,                    # Обязательно для AdaMuon
+                'lr': base_learning_rate,            # Использование обычных пробелов
+                'betas': (0.9, betta2),
+                'eps': eps,
+                'weight_decay': 0.01,
+                'nesterov': True,
+                'decouple_weight_decay': True
+            })
+            
+        # --- Группа 2: AdamW8bit (для смещений/нормализации) ---
+        if adam8bit_params:
+            print(f"Используем AdamW8bit для {len(adam8bit_params)} параметров (ndim < 2)")
+            param_groups.append({
+                'params': adam8bit_params,
+                'optimizer': 'AdamW8bit',            # Явно указываем оптимизатор
+                'use_muon': False,
+                # AdamW8bit не требует use_muon
+                'lr': base_learning_rate,            # Использование обычных пробелов
+                'betas': (0.9, betta2),
+                'eps': eps,
+                'weight_decay': 0.01,
+                'percentile_clipping': percentile_clipping
+            })
+
+        if not param_groups:
+            raise ValueError("Нет параметров для оптимизации.")
+            
+        # 3. ИНИЦИАЛИЗАЦИЯ (смешанный оптимизатор)
+        # При передаче списка групп оптимизатор AdaMuon автоматически 
+        # распознает ключи 'optimizer' и инициализирует внутренние оптимизаторы.
+        return AdaMuon(param_groups)
     else:
         raise ValueError(f"Unknown optimizer: {name}")
 
@@ -418,24 +480,22 @@ def get_negative_embedding(neg_prompt="", batch_size=1):
         seq_len = max_length
         empty_emb = torch.zeros((batch_size, seq_len, hidden_dim), dtype=dtype, device=device)
         empty_mask = torch.ones((batch_size, seq_len), dtype=torch.int64, device=device)
-        empty_pool = torch.zeros((batch_size, hidden_dim), dtype=dtype, device=device)
-        return empty_emb, empty_mask, empty_pool
+        return empty_emb, empty_mask
 
-    uncond_emb, uncond_mask, uncond_pool = encode_texts([neg_prompt])
+    uncond_emb, uncond_mask = encode_texts([neg_prompt])
     uncond_emb = uncond_emb.to(dtype=dtype, device=device).repeat(batch_size, 1, 1)
     uncond_mask = uncond_mask.to(device=device).repeat(batch_size, 1)
-    uncond_pool = uncond_pool.to(dtype=dtype, device=device).repeat(batch_size, 1)
 
-    return uncond_emb, uncond_mask, uncond_pool
+    return uncond_emb, uncond_mask
     
 # Получаем негативные (пустые) условия для валидации
-uncond_emb, uncond_mask, uncond_pool = get_negative_embedding("low quality")
+uncond_emb, uncond_mask = get_negative_embedding("low quality")
 
-# --- [UPDATED] Функция генерации семплов (принимает pool) ---
+# --- Функция генерации семплов  ---
 @torch.compiler.disable()
 @torch.no_grad()
 def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
-    uncond_emb, uncond_mask, uncond_pool = uncond_data
+    uncond_emb, uncond_mask = uncond_data
     
     original_model = None
     try:
@@ -449,13 +509,12 @@ def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
         all_generated_images = []
         all_captions = [] 
         
-        # Распаковываем 5 элементов (добавились mask и pool)
-        for size, (sample_latents, sample_text_embeddings, sample_mask, sample_pool, sample_text) in fixed_samples_cpu.items():
+        # Распаковываем 5 элементов (добавились mask)
+        for size, (sample_latents, sample_text_embeddings, sample_mask, sample_text) in fixed_samples_cpu.items():
             width, height = size
             sample_latents = sample_latents.to(dtype=dtype, device=device)
             sample_text_embeddings = sample_text_embeddings.to(dtype=dtype, device=device)
             sample_mask = sample_mask.to(device=device)
-            sample_pool = sample_pool.to(dtype=dtype, device=device)
         
             latents = torch.randn(
                 sample_latents.shape,
@@ -483,15 +542,10 @@ def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
                     neg_mask_batch = uncond_mask[0:1].expand(curr_batch_size, -1)
                     attention_mask_batch = torch.cat([neg_mask_batch, sample_mask], dim=0)
 
-                    # 3. Pools
-                    neg_pool_batch = uncond_pool[0:1].expand(curr_batch_size, -1)
-                    pooled_batch = torch.cat([neg_pool_batch, sample_pool], dim=0)
-
                 else:
                     latent_model_input = latents
                     text_embeddings_batch = sample_text_embeddings
                     attention_mask_batch = sample_mask
-                    pooled_batch = sample_pool
 
                 # Предсказание с передачей всех условий
                 model_out = original_model(
@@ -499,7 +553,6 @@ def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
                     t, 
                     encoder_hidden_states=text_embeddings_batch, 
                     encoder_attention_mask=attention_mask_batch,
-                    #added_cond_kwargs={"text_embeds": pooled_batch} # <--- ПУЛИНГ ЗДЕСЬ
                 )
                 flow = getattr(model_out, "sample", model_out)
         
@@ -560,7 +613,7 @@ def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
 if accelerator.is_main_process:
     if save_model:
         print("Генерация сэмплов до старта обучения...")
-        generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask, uncond_pool), 0)
+        generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), 0)
 accelerator.wait_for_everyone()
 
 def save_checkpoint(unet, variant=""):
@@ -601,8 +654,7 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
     accelerator.wait_for_everyone()
     unet.train()
     
-    # --- [UPDATED] Распаковка 4 элементов ---
-    for step, (latents, embeddings, attention_mask, pooled_embeddings) in enumerate(dataloader):
+    for step, (latents, embeddings, attention_mask) in enumerate(dataloader):
         with accelerator.accumulate(unet):
             if save_model == False and step == 5 :
                 used_gb = torch.cuda.max_memory_allocated() / 1024**3
@@ -617,13 +669,12 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             # делаем integer timesteps для UNet
             timesteps = (t * scheduler.config.num_train_timesteps).long()
             
-            # --- [UPDATED] Вызов UNet с маской и пулингом ---
+            # --- Вызов UNet с маской  ---
             model_pred = unet(
                 noisy_latents, 
                 timesteps, 
                 encoder_hidden_states=embeddings,
-                encoder_attention_mask=attention_mask,
-                #added_cond_kwargs={"text_embeds": pooled_embeddings} # <--- Передача пулинга
+                encoder_attention_mask=attention_mask
             ).sample
             
             target = noise - latents 
@@ -641,8 +692,8 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             grad = 0.0
             if not fbp:
                 if accelerator.sync_gradients:
-                    with torch.amp.autocast('cuda', enabled=False):
-                        grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
+                    #with torch.amp.autocast('cuda', enabled=False):
+                    grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
                     grad = float(grad_val)
                     optimizer.step()
                     lr_scheduler.step()
@@ -669,8 +720,8 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
                             comet_experiment.log_metrics(log_data, step=global_step)
     
                     if global_step % sample_interval == 0:
-                        # Передаем tuple (emb, mask, pool) для негатива
-                        generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask, uncond_pool), global_step)
+                        # Передаем tuple (emb, mask) для негатива
+                        generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
                         last_n = sample_interval
     
                         if save_model:

unet/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:987ce2ea59106bf0806b16c21fff90d6fe6e7a90101e0157d8295803dbf34f2b
-size 1892
+oid sha256:8ab7222cdd538ff5178adc870a764d22ab24a185f0a7b63852ea728b3b09fcff
+size 1876

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1924f863b340d5698a8f86471a5ac9ad0f6eda9477d1b43a1a7a146300c5bb60
-size 3302474272
+oid sha256:2b202a600c7baf31b986d98f00109d0c396f0fddc91bab63c49a01c3477eeba1
+size 6604736640