dataset_fromfolder.py

from datasets import Dataset, load_from_disk, concatenate_datasets
from diffusers import AutoencoderKL
from torchvision.transforms import Resize, ToTensor, Normalize, Compose, InterpolationMode, Lambda
from transformers import AutoModel, AutoImageProcessor, AutoTokenizer
import torch
import os
import gc
import numpy as np
from PIL import Image
from tqdm import tqdm
import random
import json
import shutil
import time
from datetime import timedelta

# ---------------- 1️⃣ Настройки ----------------
dtype = torch.float16
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 20
min_size = 640
max_size = 1152
step = 64
img_share = 0.1
empty_share = 0.1
limit = 0
# Основная процедура обработки
folder_path = "/workspace/all2"
save_path = "/workspace/1152p2"
os.makedirs(save_path, exist_ok=True)

# Функция для очистки CUDA памяти
def clear_cuda_memory():
    if torch.cuda.is_available():
        used_gb = torch.cuda.max_memory_allocated() / 1024**3
        print(f"used_gb: {used_gb:.2f} GB")
        torch.cuda.empty_cache()
        gc.collect()

# ---------------- 2️⃣ Загрузка моделей ----------------
def load_models():
    print("Загрузка моделей...")
    vae = AutoencoderKL.from_pretrained("AuraDiffusion/16ch-vae", torch_dtype=dtype).to(device).eval()
    model = AutoModel.from_pretrained("visheratin/mexma-siglip", torch_dtype=dtype, trust_remote_code=True, optimized=True).to(device)
    processor = AutoImageProcessor.from_pretrained("visheratin/mexma-siglip", use_fast=True)
    tokenizer = AutoTokenizer.from_pretrained("visheratin/mexma-siglip")
    return vae, model, processor, tokenizer

vae, model, processor, tokenizer = load_models()

# ---------------- 3️⃣ Трансформации ----------------
def get_image_transform(min_size=256, max_size=512, step=64):
    def transform(img, dry_run=False):
        # Сохраняем исходные размеры изображения
        original_width, original_height = img.size

        # 0. Ресайз: масштабируем изображение, чтобы максимальная сторона была равна max_size
        if original_width >= original_height:
            new_width = max_size
            new_height = int(max_size * original_height / original_width)
        else:
            new_height = max_size
            new_width = int(max_size * original_width / original_height)

        if new_height < min_size or new_width < min_size:
            # 1. Ресайз: масштабируем изображение, чтобы минимальная сторона была равна min_size
            if original_width <= original_height:
                new_width = min_size
                new_height = int(min_size * original_height / original_width)
            else:
                new_height = min_size
                new_width = int(min_size * original_width / original_height)

        # 2. Проверка: если одна из сторон превышает max_size, готовимся к обрезке
        crop_width = min(max_size, (new_width // step) * step)
        crop_height = min(max_size, (new_height // step) * step)

        # Убеждаемся, что размеры обрезки не меньше min_size
        crop_width = max(min_size, crop_width)
        crop_height = max(min_size, crop_height)
        
        # Если запрошен только предварительный расчёт размеров
        if dry_run:
            return crop_width, crop_height
        
        # Конвертация в RGB и ресайз
        img_resized = img.convert("RGB").resize((new_width, new_height), Image.LANCZOS)
        
        # Определение координат обрезки (обрезаем с учетом вотермарок - треть сверху)
        top = (new_height - crop_height) // 3
        left = 0
        
        # Обрезка изображения
        img_cropped = img_resized.crop((left, top, left + crop_width, top + crop_height))
        
        # Сохраняем итоговые размеры после всех преобразований
        final_width, final_height = img_cropped.size
        
        # тензор
        img_tensor = ToTensor()(img_cropped)
        img_tensor = Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])(img_tensor)
        return img_tensor, img_cropped, final_width, final_height

    return transform

# ---------------- 4️⃣ Функции обработки ----------------
def encode_images_batch(images, processor, model):
    pixel_values = torch.stack([processor(images=img, return_tensors="pt")["pixel_values"].squeeze(0) for img in images]).to(device, dtype)

    with torch.inference_mode():
        image_embeddings = model.vision_model(pixel_values).pooler_output

    return image_embeddings.unsqueeze(1).cpu().numpy()

def encode_texts_batch(texts, tokenizer, model):
    with torch.inference_mode():
        text_tokenized = tokenizer(texts, return_tensors="pt", padding="max_length",
            max_length=512,
            truncation=True).to(device)
        text_embeddings = model.encode_texts(text_tokenized.input_ids, text_tokenized.attention_mask)
    return text_embeddings.unsqueeze(1).cpu().numpy()

def maybe_empty_label(label, prob=0.01):
    return "" if random.random() < prob else label

def encode_to_latents(images, texts):
    transform = get_image_transform(min_size, max_size, step)
    try:
        # Обработка каждого изображения и сохранение размеров
        latents_list = []
        widths = []
        heights = []
        pil_images = []
        
        for img in images:
            try:
                # Применяем единую трансформацию для всех случаев
                transformed_img, pil_img, final_width, final_height = transform(img)
                pil_images.append(pil_img)  # Сохраняем преобразованное PIL изображение
                widths.append(final_width)
                heights.append(final_height)
                
                # Кодируем в VAE латенты
                img_tensor = transformed_img.unsqueeze(0).to(device, dtype)
                with torch.no_grad():
                    posterior = vae.encode(img_tensor).latent_dist.mode()
                    z = (posterior - vae.config.shift_factor) * vae.config.scaling_factor
                    latents_list.append(z.cpu().numpy())
            except Exception as e:
                print(f"Ошибка при кодировании VAE: {e}")
                continue
        
        latents = np.concatenate(latents_list, axis=0)
        
        # Выбираем между текстовыми и изображение-эмбеддингами
        if random.random() < img_share:
            # Используем уже преобразованные изображения для эмбеддингов
            embeddings = encode_images_batch(pil_images, processor, model)
        else:
            text_labels_with_empty = [maybe_empty_label(lbl, empty_share) for lbl in texts]
            embeddings = encode_texts_batch(text_labels_with_empty, tokenizer, model)
        
        return {
            "vae": latents,
            "embeddings": embeddings,
            "text": texts,
            "width": widths,
            "height": heights
        }
    except Exception as e:
        print(f"Критическая ошибка в encode_to_latents: {e}")
        raise

# ---------------- 5️⃣ Обработка папки с изображениями и текстами ----------------
def process_folder(folder_path, limit=None):
    """
    Рекурсивно обходит указанную директорию и все вложенные директории,
    собирая пути к изображениям и соответствующим текстовым файлам.
    """
    image_paths = []
    text_paths = []
    width = []
    height = []
    transform = get_image_transform(min_size, max_size, step)
    
    # Используем os.walk для рекурсивного обхода директорий
    for root, dirs, files in os.walk(folder_path):
        for filename in files:
            # Проверяем, является ли файл изображением
            if filename.lower().endswith((".jpg", ".jpeg", ".png")):
                image_path = os.path.join(root, filename)
                try:
                    img = Image.open(image_path)
                except Exception as e:
                    print(f"Ошибка при открытии {image_path}: {e}")
                    os.remove(image_path)
                    text_path = os.path.splitext(image_path)[0] + ".txt"
                    if os.path.exists(text_path):
                        os.remove(text_path)
                    continue
                # Применяем трансформацию только для получения размеров
                w, h = transform(img, dry_run=True)
                # Формируем путь к текстовому файлу
                text_path = os.path.splitext(image_path)[0] + ".txt"
                
                # Добавляем пути, если текстовый файл существует
                if os.path.exists(text_path) and min(w, h)>0:
                    image_paths.append(image_path)
                    text_paths.append(text_path)
                    width.append(w)  # Добавляем в список
                    height.append(h)  # Добавляем в список
                    
                    # Проверяем ограничение на количество
                    if limit and limit>0 and len(image_paths) >= limit:
                        print(f"Достигнут лимит в {limit} изображений")
                        return image_paths, text_paths, width, height
    
    print(f"Найдено {len(image_paths)} изображений с текстовыми описаниями")
    return image_paths, text_paths, width, height
    
def process_in_chunks(image_paths, text_paths, width, height, chunk_size=50000, batch_size=1):
    total_files = len(image_paths)
    start_time = time.time()
    chunks = range(0, total_files, chunk_size)
    
    for chunk_idx, start in enumerate(chunks, 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] if isinstance(width, list) else [width] * len(chunk_image_paths)
        chunk_heights = height[start:end] if isinstance(height, list) else [height] * len(chunk_image_paths)
        
        # Чтение текстов
        chunk_texts = []
        for text_path in chunk_text_paths:
            try:
                with open(text_path, 'r', encoding='utf-8') as f:
                    text = f.read().strip()
                chunk_texts.append(text)
            except Exception as e:
                print(f"Ошибка чтения {text_path}: {e}")
                chunk_texts.append("")
        
        # Группируем изображения по размерам
        size_groups = {}
        for i in range(len(chunk_image_paths)):
            size_key = (chunk_widths[i], chunk_heights[i])
            if size_key not in size_groups:
                size_groups[size_key] = {"image_paths": [], "texts": []}
            size_groups[size_key]["image_paths"].append(chunk_image_paths[i])
            size_groups[size_key]["texts"].append(chunk_texts[i])
        
        # Обрабатываем каждую группу размеров отдельно
        for size_key, group_data in size_groups.items():
            print(f"Обработка группы с размером {size_key[0]}x{size_key[1]} - {len(group_data['image_paths'])} изображений")
            
            group_dataset = Dataset.from_dict({
                "image_path": group_data["image_paths"],
                "text": group_data["texts"]
            })
            
            # Теперь можно использовать указанный batch_size, т.к. все изображения одного размера
            processed_group = group_dataset.map(
                lambda examples: encode_to_latents(
                    [Image.open(path) for path in examples["image_path"]],
                    examples["text"]
                ),
                batched=True,
                batch_size=batch_size,
                remove_columns=["image_path"],
                desc=f"Обработка группы размера {size_key[0]}x{size_key[1]}"
            )
            
            # Сохраняем результаты группы
            group_save_path = f"{save_path}_temp/chunk_{chunk_idx}_size_{size_key[0]}x{size_key[1]}"
            processed_group.save_to_disk(group_save_path)
            clear_cuda_memory()
            elapsed = time.time() - start_time
            processed = (chunk_idx - 1) * chunk_size + sum([len(sg["image_paths"]) for sg in list(size_groups.values())[:list(size_groups.values()).index(group_data) + 1]])
            if processed > 0:
                remaining = (elapsed / processed) * (total_files - processed)
                elapsed_str = str(timedelta(seconds=int(elapsed)))
                remaining_str = str(timedelta(seconds=int(remaining)))
                print(f"ETA: Прошло {elapsed_str}, Осталось {remaining_str}, Прогресс {processed}/{total_files} ({processed/total_files:.1%})")

# ---------------- 7️⃣ Объединение чанков ----------------
def combine_chunks(temp_path, final_path):
    """Объединение обработанных чанков в финальный датасет"""
    chunks = sorted([
        os.path.join(temp_path, d) 
        for d in os.listdir(temp_path) 
        if d.startswith("chunk_")
    ])
    
    datasets = [load_from_disk(chunk) for chunk in chunks]
    combined = concatenate_datasets(datasets)
    combined.save_to_disk(final_path)
    
    print(f"✅ Датасет успешно сохранен в: {final_path}")


# Создаем временную папку для чанков
temp_path = f"{save_path}_temp"
os.makedirs(temp_path, exist_ok=True)

# Получаем список файлов
image_paths, text_paths, width, height = process_folder(folder_path,limit)
print(f"Всего найдено {len(image_paths)} изображений")

# Обработка с чанкованием
process_in_chunks(image_paths, text_paths, width, height, chunk_size=50000, batch_size=batch_size)

# Объединение чанков в финальный датасет
combine_chunks(temp_path, save_path)

# Удаление временной папки
try:
    shutil.rmtree(temp_path)
    print(f"✅ Временная папка {temp_path} успешно удалена")
except Exception as e:
    print(f"⚠️ Ошибка при удалении временной папки: {e}")