2511

dataset.py

@@ -22,15 +22,15 @@ import tempfile
 # ---------------- 1️⃣ Настройки ----------------
 dtype = torch.float16
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-batch_size = 5
+batch_size = 10
 min_size = 320 #192 #256 #192
 max_size = 640 #384 #256 #384
 step = 64 #64
 empty_share = 0.05
 limit = 0
 # Основная процедура обработки
-folder_path = "/workspace/1" #alchemist"
-save_path = "/workspace/1_640" #"alchemist"
+folder_path = "/workspace/tar" #alchemist"
+save_path = "/workspace/640" #"alchemist"
 dir_tmp = "/workspace/tmp"
 os.makedirs(save_path, exist_ok=True)
 
@@ -355,19 +355,58 @@ def process_in_chunks(image_paths, text_paths, width, height, chunk_size=10000,
                 print(f"ETA: Прошло {elapsed_str}, Осталось {remaining_str}, Прогресс {processed}/{total_files} ({processed/total_files:.1%})")
 
 # ---------------- 7️⃣ Объединение ----------------
+def safe_repack(ds):
+    """
+    Перепаковка датасета в маленькие .arrow файлы.
+    ЧТО ДЕЛАЕТ:
+      - уменьшает Arrow chunk size
+      - writer_batch_size=1000 → Arrow 30–60 MB
+    """
+    return ds.map(lambda x: x, batched=True, writer_batch_size=1000)
+
+
 def combine_chunks(temp_path, final_path):
-    archives = sorted([os.path.join(temp_path, f) for f in os.listdir(temp_path) if f.endswith(".tar.zst")])
-    datasets = []
-    for arc in archives:
-        #tmp = tempfile.mkdtemp()
+    archives = sorted(
+        f for f in os.listdir(temp_path) 
+        if f.endswith(".tar.zst")
+    )
+    archives = [os.path.join(temp_path, f) for f in archives]
+
+    print(f"Найдено {len(archives)} архивов.")
+
+    # Инициализируем пустой датасет
+    merged_ds = None
+
+    for i, arc in enumerate(archives):
+        print(f"[{i+1}/{len(archives)}] Обрабатываю {arc}")
+
+        # Распаковка
         tmp = tempfile.mkdtemp(dir=dir_tmp)
-        os.makedirs(tmp, exist_ok=True)
         subprocess.run(["tar", "-I", "zstd", "-xf", arc, "-C", tmp], check=True)
+
+        # Загрузка датасета
         ds = load_from_disk(tmp)
-        datasets.append(ds)
+
+        # Перепаковка чанк-датасета, чтобы уменьшить Arrow-файлы
+        ds = safe_repack(ds)
+
+        # Мерж
+        if merged_ds is None:
+            merged_ds = ds
+        else:
+            merged_ds = concatenate_datasets([merged_ds, ds])
+
+        # cleanup
         shutil.rmtree(tmp)
-    combined = concatenate_datasets(datasets)
-    combined.save_to_disk(final_path)
+        os.remove(arc)
+
+    # Финальная перепаковка
+    print("⚙️ Финальная перепаковка...")
+    merged_ds = safe_repack(merged_ds)
+
+    print("💾 Финальное сохранение...")
+    merged_ds.save_to_disk(final_path)
+
     print(f"✅ Датасет сохранён: {final_path}")
 
 # Создаем временную папку для чанков

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f7226baa7a6111c3ace4e7692400c546296e01273c4d946c1d13753454c2d359
+oid sha256:a986a5cef456e178da40d574d94ade01ee5d9489d1540f318bcf9f696fc658d3
 size 6184944280