vae8x16x32ch

dataset.py

@@ -9,34 +9,32 @@ import shutil
 import time
 
 from datasets import Dataset, load_from_disk, concatenate_datasets
-from diffusers import AutoencoderKL,AutoencoderKLWan,AsymmetricAutoencoderKL,AutoencoderKLFlux2
+from diffusers import AsymmetricAutoencoderKL
 from torchvision.transforms import Resize, ToTensor, Normalize, Compose, InterpolationMode, Lambda
 from transformers import AutoModel, AutoImageProcessor, AutoTokenizer, AutoModelForCausalLM
 from typing import Dict, List, Tuple, Optional, Any
 from PIL import Image
 from tqdm import tqdm
 from datetime import timedelta
-
-# --- NEW: accelerate ---
 from accelerate import Accelerator
+
 accelerator = Accelerator()
 device = accelerator.device
 is_main_process = accelerator.is_main_process
 process_index = accelerator.process_index
 num_processes = accelerator.num_processes
-# --- END NEW ---
 
 # ---------------- 1️⃣ Настройки ----------------
 dtype = torch.float16
-batch_size = 5
-min_size = 640
-max_size = 1280
-step = 64
+batch_size = 4
+min_size = 384
+max_size = 704
+step = 32
 empty_share = 0.0
 limit = 0
 
-folder_path = "/workspace/ds"
-save_path = "/workspace/ds1234_flux32"
+folder_path = "/workspace/sdxs-1b/datasets/ds1234_noanime_1280"
+save_path = "/workspace/sdxs-1b/datasets/ds1234_noanime_704_vae8x16x"
 os.makedirs(save_path, exist_ok=True)
 
 def clear_cuda_memory():
@@ -49,7 +47,7 @@ def clear_cuda_memory():
 # ---------------- 2️⃣ Загрузка моделей ----------------
 def load_models():
     print(f"[GPU {process_index}] Загрузка моделей...")
-    vae = AutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
+    vae = AsymmetricAutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
     return vae
 
 vae = load_models()
@@ -197,7 +195,102 @@ def process_folder(folder_path, limit=None):
     print(f"Найдено {len(image_paths)} изображений")
     return image_paths, text_paths, width, height
 
-def process_in_chunks(image_paths, text_paths, width, height, chunk_size=10000, batch_size=1):
+def process_in_chunks(image_paths, text_paths, width, height, chunk_size=10000, batch_size=4):
+    total_files = len(image_paths)
+
+    for chunk_idx, start in enumerate(range(0, total_files, chunk_size), 1):
+        end = min(start + chunk_size, total_files)
+
+        chunk_image_paths = image_paths[start:end]
+        chunk_text_paths = text_paths[start:end]
+        chunk_widths = width[start:end]
+        chunk_heights = height[start:end]
+
+        # --- читаем тексты ---
+        chunk_texts = []
+        for text_path in chunk_text_paths:
+            try:
+                with open(text_path, "r", encoding="utf-8") as f:
+                    chunk_texts.append(f.read().strip())
+            except:
+                chunk_texts.append("")
+
+        # --- группировка по размеру ---
+        size_groups = {}
+        for i in range(len(chunk_image_paths)):
+            key = (chunk_widths[i], chunk_heights[i])
+            if key not in size_groups:
+                size_groups[key] = {"image_paths": [], "texts": []}
+            size_groups[key]["image_paths"].append(chunk_image_paths[i])
+            size_groups[key]["texts"].append(chunk_texts[i])
+
+        # --- обработка групп ---
+        for size_key, group_data in size_groups.items():
+            image_list = group_data["image_paths"]
+            text_list = group_data["texts"]
+
+            latents_all = []
+            texts_all = []
+            widths_all = []
+            heights_all = []
+
+            for i in range(0, len(image_list), batch_size):
+                batch_paths = image_list[i:i + batch_size]
+                batch_texts = text_list[i:i + batch_size]
+
+                batch_imgs = []
+
+                for p in batch_paths:
+                    try:
+                        with Image.open(p) as img:
+                            img = img.convert("RGB")
+                            batch_imgs.append(img.copy())
+                    except Exception as e:
+                        print(f"[GPU {process_index}] Ошибка загрузки: {p} | {e}")
+
+                if len(batch_imgs) == 0:
+                    continue
+
+                try:
+                    out = encode_to_latents(batch_imgs, batch_texts)
+                except Exception as e:
+                    print(f"[GPU {process_index}] Ошибка encode: {e}")
+                    continue
+
+                if out is None:
+                    continue
+
+                latents_all.extend(out["vae"])
+                texts_all.extend(out["text"])
+                widths_all.extend(out["width"])
+                heights_all.extend(out["height"])
+
+                # чуть чистим память
+                del batch_imgs, out
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+
+            if len(latents_all) == 0:
+                continue
+
+            # --- сохраняем ---
+            group_save_path = f"{save_path}_temp/chunk_{chunk_idx}_{size_key[0]}x{size_key[1]}_proc_{process_index}"
+
+            dataset_dict = {
+                "vae": latents_all,
+                "text": texts_all,
+                "width": widths_all,
+                "height": heights_all,
+            }
+
+            ds = Dataset.from_dict(dataset_dict)
+            ds.save_to_disk(group_save_path)
+
+            print(f"[GPU {process_index}] Saved: {group_save_path}")
+
+            clear_cuda_memory()
+            
+def process_in_chunks2(image_paths, text_paths, width, height, chunk_size=10000, batch_size=1):
     total_files = len(image_paths)
     start_time = time.time()
 
@@ -237,7 +330,7 @@ def process_in_chunks(image_paths, text_paths, width, height, chunk_size=10000,
             )
 
             # --- NEW: уникальный путь ---
-            group_save_path = f"{save_path}_temp/proc_{process_index}_chunk_{chunk_idx}_{size_key[0]}x{size_key[1]}"
+            group_save_path = f"{save_path}_temp/chunk_{chunk_idx}_{size_key[0]}x{size_key[1]}_proc_{process_index}_"
             # --- END NEW ---
 
             processed_group.save_to_disk(group_save_path)
@@ -270,7 +363,7 @@ text_paths = [text_paths[i] for i in sorted_indices]
 width = [width[i] for i in sorted_indices]
 height = [height[i] for i in sorted_indices]
 
-# --- NEW: shard по GPU ---
+# --- shard по GPU ---
 indices = list(range(len(image_paths)))
 indices = indices[process_index::num_processes]
 
@@ -280,13 +373,10 @@ width = [width[i] for i in indices]
 height = [height[i] for i in indices]
 
 print(f"[GPU {process_index}] обрабатывает {len(image_paths)} файлов")
-# --- END NEW ---
 
 process_in_chunks(image_paths, text_paths, width, height, chunk_size=20000, batch_size=batch_size)
 
-# --- NEW: sync ---
 accelerator.wait_for_everyone()
-# --- END NEW ---
 
 # --- NEW: только главный процесс ---
 if is_main_process:
@@ -300,5 +390,4 @@ if is_main_process:
     try:
         shutil.rmtree(temp_path)
     except:
-        pass
-# --- END NEW ---
+        pass

model_index.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f1b6ce208190165975458f502eb3e7ad6d4a5dd54a507f8dd727e636e363ae93
-size 428
+oid sha256:b9fe0891c1d3f4f0b2a8cbca077be3533f28306768c3ea8d5256924fc677a4b1
+size 438

src/test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ff518f6f2623a12e096a9b277d623c0ba5324e6a54d3383ddd5cd85ace78bf47
-size 5718386
+oid sha256:dba3b9d70d28c3827ea6f79c9f70740f79ae7eec0f37291467c4e1a9d708f310
+size 6906512

test.ipynb

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

train.py

@@ -11,7 +11,7 @@ from torch.utils.data import DataLoader, Sampler
 from torch.utils.data.distributed import DistributedSampler
 from torch.optim.lr_scheduler import LambdaLR
 from collections import defaultdict
-from diffusers import UNet2DConditionModel,AutoencoderKL,FlowMatchEulerDiscreteScheduler
+from diffusers import UNet2DConditionModel,AsymmetricAutoencoderKL,FlowMatchEulerDiscreteScheduler
 from accelerate import Accelerator, DeepSpeedPlugin
 from datasets import load_from_disk
 from tqdm import tqdm
@@ -31,7 +31,7 @@ from transformers import Qwen3_5Tokenizer, Qwen3_5ForConditionalGeneration
 import argparse
 
 # --------------------------- Параметры ---------------------------
-ds_path = "datasets/ds1234_1280"
+ds_path = "datasets/ds1234_noanime_704_vae8x16x"
 project = "unet"
 # 1. Считаем локальный батч для ОДНОЙ карты (3 на каждые 32 Гб)
 gpu_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
@@ -48,8 +48,8 @@ sample_interval_share = 20
 cfg_dropout = 0.10
 max_length = 248
 use_wandb = False
-use_comet_ml = True
-save_model = True
+use_comet_ml = False
+save_model = False
 use_decay = True
 fbp = False
 optimizer_type = "adam8bit"
@@ -148,7 +148,7 @@ if accelerator.is_main_process:
         comet_experiment.log_parameters(hyper_params)
 
 # --------------------------- Загрузка моделей ---------------------------
-vae = AutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
+vae = AsymmetricAutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
 tokenizer = Qwen3_5Tokenizer.from_pretrained("tokenizer")
 text_encoder = Qwen3_5ForConditionalGeneration.from_pretrained("text_encoder", torch_dtype=torch.float16).to(device).eval()
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("scheduler")
@@ -432,14 +432,8 @@ if limit > 0:
 else:
     dataset = load_from_disk(ds_path)
 
-dataset = dataset.filter(
-    lambda x: [not (path.startswith("//workspace/ds/animesfw") or path.startswith("//workspace/animesfw") or path.startswith("//workspace/ds/d4/animesfw"))  for path in x["image_path"]],
 
-    batched=True,
-    batch_size=10000, # обрабатываем по 10к строк за раз
-    num_proc=8
-)
-print(f"Осталось примеров после фильтрации: {len(dataset)}")
+print(f"images: {len(dataset)}")
 
 # --- Collate Function ---
 def collate_fn_simple(batch):

vae/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a35c34ddc8de51cec98075acd06feb9c55fd79a75d4f7e055adb5f4809e637c3
-size 2331
+oid sha256:26f99e8253949bc6a406c5f7c7f1358c76d20eb3564e579378c21e0596c466cf
+size 2379

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6369e370ff02168a240a9ebfd47810dd7babb36f76b7d9999e5d78cb4a1976c2
-size 336212308
+oid sha256:780993a422cf0d783088c5116063ab6e3b1453a5256f205b78b1c67617f17b33
+size 383499124