2512

pipeline_sdxs.py

@@ -12,7 +12,7 @@ class SdxsPipelineOutput(BaseOutput):
     images: Union[List[Image.Image], np.ndarray]
 
 class SdxsPipeline(DiffusionPipeline):
-    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler, max_length: int = 150):
+    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler, max_length: int = 192):
         super().__init__()
         self.register_modules(
             vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, 
@@ -33,7 +33,7 @@ class SdxsPipeline(DiffusionPipeline):
     
         # Если промпты не заданы, используем пустые эмбеддинги
         if prompt is None and negative_prompt is None:
-            hidden_dim = 1024  # Размерность эмбеддинга Qwen3-0.6B
+            hidden_dim = 1024  # Размерность эмбеддинга
             seq_len = self.max_length
             batch_size = 1
             # ИЗМЕНЕНО: Возвращаем три элемента: embeds, mask, pooled
@@ -42,7 +42,7 @@ class SdxsPipeline(DiffusionPipeline):
             empty_pooled = torch.zeros((batch_size, hidden_dim), dtype=dtype, device=device)
             return empty_embeds, empty_mask, empty_pooled
     
-        # Токенизация с фиксированным max_length=150 и padding="max_length"
+        # Токенизация с фиксированным max_length и padding="max_length"
         def encode_texts(texts, max_length=self.max_length):
             with torch.no_grad():
                 if isinstance(texts, str):
@@ -70,7 +70,7 @@ class SdxsPipeline(DiffusionPipeline):
                 outs = self.text_encoder(**toks, output_hidden_states=True, return_dict=True)
                 
                 # Токен-эмбеддинги (для Cross-Attention)
-                hidden = outs.hidden_states[-1] # Используем last hidden state -2???
+                hidden = outs.hidden_states[-2] # Используем last hidden state -2???
                 # Маска внимания (для Cross-Attention)
                 attention_mask = toks["attention_mask"] 
                 

src/conv2048.py

@@ -0,0 +1,97 @@
+import torch
+from diffusers import UNet2DConditionModel
+import os
+from safetensors.torch import load_file as safe_load
+
+# --- КОНСТАНТЫ ---
+OLD_UNET_PATH = "/workspace/sdxs/unet_old" 
+NEW_UNET_PATH = "/workspace/sdxs/unet"     
+
+def transfer_unet_weights_fix(old_path: str, new_path: str):
+    
+    print(f"Загрузка новой UNet (Конфиг 2048D) из: {new_path}")
+    new_unet = UNet2DConditionModel.from_pretrained(new_path, low_cpu_mem_usage=False)
+
+    # 2. Прямая загрузка state dict старой модели (с обработкой формата)
+    print(f"Прямая загрузка весов старой UNet из: {old_path}")
+    
+    weights_file = next((f for f in os.listdir(old_path) if f.endswith(('.safetensors', '.bin'))), None)
+    if not weights_file:
+        print("Ошибка: не найден файл весов (.safetensors или .bin) в старой папке.")
+        return None
+    
+    if weights_file.endswith('.safetensors'):
+        print(f"Обнаружен файл Safetensors ({weights_file}). Использую safe_load.")
+        old_state_dict = safe_load(f"{old_path}/{weights_file}")
+    elif weights_file.endswith('.bin'):
+        print(f"Обнаружен файл PyTorch (.bin) ({weights_file}). Использую torch.load с weights_only=False.")
+        old_state_dict = torch.load(
+            f"{old_path}/{weights_file}", 
+            map_location='cpu',
+            weights_only=False 
+        )
+    else:
+        print(f"Ошибка: Не удалось загрузить файл {weights_file}. Проверьте формат.")
+        return None
+    
+    if "state_dict" in old_state_dict:
+        old_state_dict = old_state_dict["state_dict"]
+    
+    # 3. Перенос весов с обработкой RuntimeError
+    print("Начало переноса весов с пропуском несовместимых слоев...")
+    
+    # --- НОВОЕ: Предварительный подсчет совпадающих ключей ---
+    total_keys = len(old_state_dict)
+    matching_keys = 0
+    size_mismatch_keys = 0
+    
+    # Используем ключи новой модели для сравнения, чтобы убедиться, 
+    # что ключи, которые отсутствуют в старой, не учитываются в "перенесенных".
+    new_keys = new_unet.state_dict().keys()
+
+    for name, old_param in old_state_dict.items():
+        if name in new_keys:
+            new_param = new_unet.state_dict()[name]
+            # Считаем только те, где имя и размер совпали
+            if old_param.shape == new_param.shape:
+                matching_keys += 1
+            # Считаем те, где имя совпало, но размер изменился (mismatch)
+            else:
+                size_mismatch_keys += 1
+        # Ключи, которые есть в старой, но нет в новой (unexpected), 
+        # также будут пропущены. Но нас в первую очередь интересуют size mismatch.
+
+    # Пытаемся загрузить веса, зная, что это вызовет RuntimeError
+    try:
+        # Запуск actual переноса. Ключи будут перенесены.
+        new_unet.load_state_dict(old_state_dict, strict=False)
+        
+    except RuntimeError as e:
+        # Это ожидаемый блок! Он ловит RuntimeError, вызванный mismatch size.
+        # Веса, которые совпали, УЖЕ перенесены в new_unet.
+        
+        print("\n--- Отчет о переносе весов ---")
+        print("⚠️ Обнаружен ожидаемый **RuntimeError** из-за несовпадения размеров. Это нормально!")
+        
+        print(f"💡 **УСПЕШНО перенесенных ключей (совпадающий размер): {matching_keys} шт.**")
+        print(f"❌ **Пропущенных ключей (несовпадение размера): {size_mismatch_keys} шт.**")
+        
+        # Мы всё равно сохраняем UNet, так как большая часть весов перенесена
+        new_unet.save_pretrained(new_path)
+        print(f"\n✅ Новая UNet (с перенесенными весами) сохранена по пути: {new_path}")
+        return new_unet
+        
+    # Блок на случай, если чудом ошибки не возникло (для полноты)
+    print("\n--- Отчет о переносе весов ---")
+    print(f"✅ Успешно перенесены совпадающие веса (основная часть UNet).")
+    print(f"💡 **УСПЕШНО перенесенных ключей (совпадающий размер): {matching_keys} шт.**")
+    print(f"❌ **Пропущенных ключей (несовпадение размера): {size_mismatch_keys} шт.**")
+    #print(f"❌ Слои, требующие переобучения (измененный размер): {len(incompatible_keys.unexpected_keys)}")
+    
+    new_unet.save_pretrained(new_path)
+    print(f"\n✅ Новая UNet сохранена по пути: {new_path}")
+    
+    return new_unet
+
+# --- ВЫПОЛНЕНИЕ ---
+transferred_unet = transfer_unet_weights_fix(OLD_UNET_PATH, NEW_UNET_PATH)

src/sdxs_create.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7fc8d44920e51edca4d0ab57c808991908139fb28987fb3bbfcb116068d4dc20
-size 7417
+oid sha256:d542ea503b79500cf1ee2bec8f9a82807d0520579664ab363b4611c1971620c9
+size 8018

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dcec4623b1c01ec0ea887756e54c2926705d4a18a812870e96612c672af48065
-size 5017634
+oid sha256:70b3ec42fbbada7f67f0477e673da4b15da141fa481ff04f709680ca0eaf773f
+size 4393987

text_encoder/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:39738c799f70b98f75dace818f14bfffb0b34fad25d41b80b6236e20a731c740
-size 1359
+oid sha256:b6026cadca04402d27495be77ff01ca4e1b418ea59047ad7079bf3bfd3517f11
+size 1354

text_encoder/model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a3184342876e1e64cd98164dc01aaaf69177a9d0ed80725816e3225977ab2982
-size 1192134784
+oid sha256:2a86278726ecbda59caff65f49ec07c3ec5672b044ed523780990a8c08a87fd4
+size 1192133232

train.py

@@ -26,18 +26,19 @@ from collections import deque
 from transformers import AutoTokenizer, AutoModel
 
 # --------------------------- Параметры ---------------------------
-ds_path = "/workspace/sdxs/datasets/ds1234_640"
+ds_path = "/workspace/sdxs/datasets/640"
 project = "unet"
-batch_size = 48
+batch_size = 256
 base_learning_rate = 3e-5
 min_learning_rate = 2.5e-5
-num_epochs = 5
-sample_interval_share = 1
+num_epochs = 10
+sample_interval_share = 20
+max_length = 192
 use_wandb = True
 use_comet_ml = False
 save_model = True
 use_decay = True
-fbp = False 
+fbp = False
 optimizer_type = "adam8bit"
 torch_compile = False
 unet_gradient = True
@@ -50,14 +51,14 @@ torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.enable_mem_efficient_sdp(False)
 dtype = torch.float32
 save_barrier = 1.004
-warmup_percent = 0.005
+warmup_percent = 0.01
 percentile_clipping = 98
 betta2 = 0.998
 eps = 1e-6
 clip_grad_norm = 1.0
 limit = 0
 checkpoints_folder = ""
-mixed_precision = "no" 
+mixed_precision = "bf16" 
 gradient_accumulation_steps = 1
 
 accelerator = Accelerator(
@@ -124,7 +125,7 @@ tokenizer = AutoTokenizer.from_pretrained("tokenizer")
 text_model = AutoModel.from_pretrained("text_encoder").to(device).eval()
 
 # --- [UPDATED] Функция кодирования текста (с маской и пулингом) ---
-def encode_texts(texts, max_length=150):
+def encode_texts(texts, max_length=max_length):
     # Если тексты пустые (для unconditional), создаем заглушки
     if texts is None: 
         # В случае None возвращаем нули (логика для get_negative_embedding)
@@ -158,9 +159,8 @@ def encode_texts(texts, max_length=150):
         
         outs = text_model(**toks, output_hidden_states=True, return_dict=True)
         
-        # 1. Токен-эмбеддинги (Last Hidden State - нормализованный)
-        # Используем last_hidden_state или hidden_states[-1] (если Qwen, лучше last_hidden_state)
-        hidden = outs.last_hidden_state 
+        # Используем last_hidden_state или hidden_states[-1] (если Qwen, лучше last_hidden_state - прим человека: ХУЙ)
+        hidden = outs.hidden_states[-2] 
         
         # 2. Маска внимания
         attention_mask = toks["attention_mask"] 
@@ -398,7 +398,7 @@ fixed_samples = get_fixed_samples_by_resolution(dataset)
 def get_negative_embedding(neg_prompt="", batch_size=1):
     if not neg_prompt:
         hidden_dim = 1024 
-        seq_len = 150
+        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)

unet/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9437f6d40639f1ad6d95d6586009d7b1bf0c8e99959d29da5c9c9645cae39ea3
-size 1899
+oid sha256:ae4272f7e52480762c228a2fbe1db5f361d7a5971c3855b483999fb3df2d722b
+size 1885

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:22402fe3c630ab551173c2bd11af4a7ae4ed76c04f001535f3791cd82da93719
-size 3103078992
+oid sha256:ccfa16a08b6e507835636c048b8c721186ef6832fec3ca889a7139dcd53676cf
+size 6625750656

text_encoder/generation_config.json → unet_old/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d3e057bbca66b92f33a8bdc6a1301014e0e4ab69b3b3fd2e442d9fe0c69f3431
-size 214
+oid sha256:9437f6d40639f1ad6d95d6586009d7b1bf0c8e99959d29da5c9c9645cae39ea3
+size 1899

unet_old/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22402fe3c630ab551173c2bd11af4a7ae4ed76c04f001535f3791cd82da93719
+size 3103078992