2601

sdxs_08b/train.py

@@ -1,798 +0,0 @@
-#from comet_ml import Experiment
-import os
-import math
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-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,AutoencoderKLFlux2
-from accelerate import Accelerator
-from datasets import load_from_disk
-from tqdm import tqdm
-from PIL import Image, ImageOps
-import wandb
-import random
-import gc
-from accelerate.state import DistributedType
-from torch.distributed import broadcast_object_list
-from torch.utils.checkpoint import checkpoint
-from diffusers.models.attention_processor import AttnProcessor2_0
-from datetime import datetime
-import bitsandbytes as bnb
-import torch.nn.functional as F
-from collections import deque
-from transformers import AutoTokenizer, AutoModel
-
-# --------------------------- Параметры ---------------------------
-ds_path = "/workspace/sdxs/datasets/mjnj"
-project = "sdxs_08b"
-batch_size = 128
-base_learning_rate = 4e-5 #2.7e-5
-min_learning_rate = 1e-5 #2.7e-5
-num_epochs = 15
-sample_interval_share = 3
-cfg_dropout = 0.25
-max_length = 192
-use_wandb = True
-use_comet_ml = False
-save_model = True
-use_decay = True
-fbp = False
-optimizer_type = "adam8bit"
-torch_compile = False
-unet_gradient = True
-fixed_seed = False
-shuffle = True
-comet_ml_api_key = "Agctp26mbqnoYrrlvQuKSTk6r" 
-comet_ml_workspace = "recoilme" 
-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.01
-warmup_percent = 0.01
-percentile_clipping = 95 #96 #97
-betta2 = 0.995
-eps = 1e-7
-clip_grad_norm = 1.0
-limit = 0
-checkpoints_folder = ""
-mixed_precision = "no" 
-gradient_accumulation_steps = 1
-
-accelerator = Accelerator(
-    mixed_precision=mixed_precision,
-    gradient_accumulation_steps=gradient_accumulation_steps
-)
-device = accelerator.device
-
-# Параметры для диффузии
-n_diffusion_steps = 40
-samples_to_generate = 12
-guidance_scale = 4
-
-# Папки для сохранения результатов
-generated_folder = "samples"
-os.makedirs(generated_folder, exist_ok=True)
-
-# Настройка seed
-current_date = datetime.now()
-seed = int(current_date.strftime("%Y%m%d"))
-if fixed_seed:
-    torch.manual_seed(seed)
-    np.random.seed(seed)
-    random.seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-
-# --------------------------- Параметры LoRA ---------------------------
-lora_name = "" 
-lora_rank = 32
-lora_alpha = 64
-
-print("init")
-
-loss_ratios = {
-    "mse":   1.5,
-    "mae":   0.5,
-}
-median_coeff_steps = 256 
-
-# Нормализация лоссов по медианам: считаем КОЭФФИЦИЕНТЫ
-class MedianLossNormalizer:
-    def __init__(self, desired_ratios: dict, window_steps: int):
-        # нормируем доли на случай, если сумма != 1
-        #s = sum(desired_ratios.values())
-        #self.ratios = {k: (v / s) for k, v in desired_ratios.items()}
-        self.ratios = {k: float(v) for k, v in desired_ratios.items()}
-        self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}
-        self.window = window_steps
-
-    def update_and_total(self, losses: dict):
-        """
-        losses: dict ключ->тензор (значения лоссов)
-        Поведение:
-          - буферим ABS(l) только для активных (ratio>0) лоссов
-          - coeff = ratio / median(abs(loss))
-          - total = sum(coeff * loss) по активным лоссам
-        CHANGED: буферим abs() — чтобы медиана была положительной и не ломала деление.
-        """
-        # буферим только активные лоссы
-        for k, v in losses.items():
-            if k in self.buffers and self.ratios.get(k, 0) > 0:
-                val = v.detach().abs().mean().cpu().item() # .item() лучше float() для тензоров
-                self.buffers[k].append(val)
-                #self.buffers[k].append(float(v.detach().abs().cpu()))
-
-        meds = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
-        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
-
-        # суммируем только по активным (ratio>0)
-        total = sum(coeffs[k] * losses[k] for k in coeffs if self.ratios.get(k, 0) > 0)
-        return total, coeffs, meds
-
-# создаём normalizer после определения loss_ratios
-normalizer = MedianLossNormalizer(loss_ratios, median_coeff_steps)
-
-# --------------------------- Инициализация WandB ---------------------------
-if accelerator.is_main_process:
-    if use_wandb:
-        wandb.init(project=project+lora_name, config={
-            "batch_size": batch_size,
-            "base_learning_rate": base_learning_rate,
-            "num_epochs": num_epochs,
-            "optimizer_type": optimizer_type,
-        })
-    if use_comet_ml:
-        from comet_ml import Experiment
-        comet_experiment = Experiment(
-            api_key=comet_ml_api_key,
-            project_name=project,
-            workspace=comet_ml_workspace
-        )
-        hyper_params = {
-            "batch_size": batch_size,
-            "base_learning_rate": base_learning_rate,
-            "num_epochs": num_epochs,
-        }
-        comet_experiment.log_parameters(hyper_params)
-
-# Включение Flash Attention 2/SDPA
-torch.backends.cuda.enable_flash_sdp(True)
-
-# --------------------------- Загрузка моделей ---------------------------
-vae = AutoencoderKL.from_pretrained("vae1x", torch_dtype=dtype).to("cpu").eval()
-#vae = AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-dev",subfolder="vae",torch_dtype=dtype).to(device).eval()
-tokenizer = AutoTokenizer.from_pretrained("tokenizer")
-text_model = AutoModel.from_pretrained("text_encoder").to(device).eval()
-
-# --- [UPDATED] Функция кодирования текста (с маской и пулингом) ---
-def encode_texts(texts, max_length=max_length):
-    # Если тексты пустые (для unconditional), создаем заглушки
-    if texts is None: 
-        # В случае None возвращаем нули (логика для get_negative_embedding)
-        # Но здесь мы обычно ожидаем список строк.
-        pass
-
-    with torch.no_grad():
-        if isinstance(texts, str):
-            texts = [texts]
-
-        for i, prompt_item in enumerate(texts):
-            messages = [
-                {"role": "user", "content": prompt_item},
-            ]
-            prompt_item = tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True,
-                #enable_thinking=True,
-            )
-            #print(prompt_item+"\n")
-            texts[i] = prompt_item
-            
-        toks = tokenizer(
-            texts,
-            return_tensors="pt",
-            padding="max_length",
-            truncation=True,
-            max_length=max_length
-        ).to(device)
-        
-        outs = text_model(**toks, output_hidden_states=True, return_dict=True)
-        
-        # Используем last_hidden_state или hidden_states[-1] (если Qwen, лучше last_hidden_state - прим человека: ХУЙ)
-        hidden = outs.hidden_states[-2] 
-        
-        # 2. Маска внимания
-        attention_mask = toks["attention_mask"] 
-        
-        # 3. Пулинг-эмбеддинг (Последний токен)
-        sequence_lengths = attention_mask.sum(dim=1) - 1
-        batch_size = hidden.shape[0]
-        pooled = hidden[torch.arange(batch_size, device=hidden.device), sequence_lengths]
-
-        #return hidden, attention_mask
-        # --- НОВАЯ ЛОГИКА: ОБЪЕДИНЕНИЕ ДЛЯ КРОСС-ВНИМАНИЯ ---
-        # 1. Расширяем пулинг-вектор до последовательности [B, 1, emb]
-        pooled_expanded = pooled.unsqueeze(1) 
-                    
-        # 2. Объединяем последовательность токенов и пулинг-вектор
-        # !!! ИЗМЕНЕНИЕ ЗДЕСЬ !!!: Пулинг идет ПЕРВЫМ
-        # Теперь: [B, 1 + L, emb]. Пулинг стал токеном в НАЧАЛЕ.
-        new_encoder_hidden_states = torch.cat([pooled_expanded, hidden], dim=1) 
-                    
-        # 3. Обновляем маску внимания для нового токена
-        # Маска внимания: [B, 1 + L]. Добавляем 1 в НАЧАЛО.
-        # 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
-
-shift_factor = getattr(vae.config, "shift_factor", 0.0)
-if shift_factor is None: shift_factor = 0.0
-scaling_factor = getattr(vae.config, "scaling_factor", 1.0)
-if scaling_factor is None: scaling_factor = 1.0
-
-from diffusers import FlowMatchEulerDiscreteScheduler
-num_train_timesteps = 1000
-scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=num_train_timesteps)
-
-class DistributedResolutionBatchSampler(Sampler):
-    def __init__(self, dataset, batch_size, num_replicas, rank, shuffle=True, drop_last=True):
-        self.dataset = dataset
-        self.batch_size = max(1, batch_size // num_replicas)
-        self.num_replicas = num_replicas
-        self.rank = rank
-        self.shuffle = shuffle
-        self.drop_last = drop_last
-        self.epoch = 0
-        
-        try:
-            widths = np.array(dataset["width"])
-            heights = np.array(dataset["height"])
-        except KeyError:
-            widths = np.zeros(len(dataset))
-            heights = np.zeros(len(dataset))
-        
-        self.size_keys = np.unique(np.stack([widths, heights], axis=1), axis=0)
-        self.size_groups = {}
-        for w, h in self.size_keys:
-            mask = (widths == w) & (heights == h)
-            self.size_groups[(w, h)] = np.where(mask)[0]
-            
-        self.group_num_batches = {}
-        total_batches = 0
-        for size, indices in self.size_groups.items():
-            num_full_batches = len(indices) // (self.batch_size * self.num_replicas)
-            self.group_num_batches[size] = num_full_batches
-            total_batches += num_full_batches
-            
-        self.num_batches = (total_batches // self.num_replicas) * self.num_replicas
-        
-    def __iter__(self):
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        all_batches = []
-        rng = np.random.RandomState(self.epoch)
-        
-        for size, indices in self.size_groups.items():
-            indices = indices.copy()
-            if self.shuffle:
-                rng.shuffle(indices)
-            num_full_batches = self.group_num_batches[size]
-            if num_full_batches == 0:
-                continue
-            valid_indices = indices[:num_full_batches * self.batch_size * self.num_replicas]
-            batches = valid_indices.reshape(-1, self.batch_size * self.num_replicas)
-            start_idx = self.rank * self.batch_size
-            end_idx = start_idx + self.batch_size
-            gpu_batches = batches[:, start_idx:end_idx]
-            all_batches.extend(gpu_batches)
-        
-        if self.shuffle:
-            rng.shuffle(all_batches)
-        accelerator.wait_for_everyone()
-        return iter(all_batches)
-
-    def __len__(self):
-        return self.num_batches 
-
-    def set_epoch(self, epoch):
-        self.epoch = epoch
-
-# --- [UPDATED] Функция для фиксированных семплов ---
-def get_fixed_samples_by_resolution(dataset, samples_per_group=1):
-    size_groups = defaultdict(list)
-    try:
-        widths = dataset["width"]
-        heights = dataset["height"]
-    except KeyError:
-        widths = [0] * len(dataset)
-        heights = [0] * len(dataset)
-    for i, (w, h) in enumerate(zip(widths, heights)):
-        size = (w, h)
-        size_groups[size].append(i)
-    
-    fixed_samples = {}
-    for size, indices in size_groups.items():
-        n_samples = min(samples_per_group, len(indices))
-        if len(size_groups)==1:
-            n_samples = samples_to_generate
-        if n_samples == 0:
-            continue
-        sample_indices = random.sample(indices, n_samples)
-        samples_data = [dataset[idx] for idx in sample_indices]
-        
-        latents = torch.tensor(np.array([item["vae"] for item in samples_data])).to(device=device, dtype=dtype)
-        texts = [item["text"] for item in samples_data]
-        
-        # Кодируем тексты на лету, чтобы получить маски и пулинг
-        embeddings, masks = encode_texts(texts)
-        
-        fixed_samples[size] = (latents, embeddings, masks, texts)
-    
-    print(f"Создано {len(fixed_samples)} групп фиксированных семплов по разрешениям")
-    return fixed_samples
-
-if limit > 0:
-    dataset = load_from_disk(ds_path).select(range(limit))
-else:
-    dataset = load_from_disk(ds_path)
-
-dataset = dataset.filter(
-    lambda x: [not (path.startswith("/workspace/ds/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)}")
-
-# --- [UPDATED] Collate Function ---
-def collate_fn_simple(batch):
-    # 1. Латенты (VAE)
-    latents = torch.tensor(np.array([item["vae"] for item in batch])).to(device, dtype=dtype)
-    
-    # 2. Текст берем сырой из датасета
-    raw_texts = [item["text"] for item in batch]
-    texts = [
-        "" if t.lower().startswith("zero") 
-        else "" if random.random() < cfg_dropout
-        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)
-    embeddings, attention_mask = encode_texts(texts)
-    
-    # attention_mask от токенизатора уже имеет нужный формат, но на всякий случай приведем к long
-    attention_mask = attention_mask.to(dtype=torch.int64)
-
-    return latents, embeddings, attention_mask
-
-batch_sampler = DistributedResolutionBatchSampler(
-        dataset=dataset,
-        batch_size=batch_size,
-        num_replicas=accelerator.num_processes,
-        rank=accelerator.process_index,
-        shuffle=shuffle
-    )
-
-dataloader = DataLoader(dataset, batch_sampler=batch_sampler, collate_fn=collate_fn_simple)
-if accelerator.is_main_process:
-    print("Total samples", len(dataloader))
-dataloader = accelerator.prepare(dataloader)
-
-start_epoch = 0
-global_step = 0
-total_training_steps = (len(dataloader) * num_epochs)
-world_size = accelerator.state.num_processes
-
-# Загрузка UNet
-latest_checkpoint = os.path.join(checkpoints_folder, project)
-if os.path.isdir(latest_checkpoint):
-    print("Загружаем UNet из чекпоинта:", latest_checkpoint)
-    unet = UNet2DConditionModel.from_pretrained(latest_checkpoint).to(device=device, dtype=dtype)
-    if unet_gradient:
-        unet.enable_gradient_checkpointing()
-    unet.set_use_memory_efficient_attention_xformers(False)
-    try:
-        unet.set_attn_processor(AttnProcessor2_0())
-    except Exception as e:
-        print(f"Ошибка при включении SDPA: {e}")
-        unet.set_use_memory_efficient_attention_xformers(True)
-else:
-    raise FileNotFoundError(f"UNet checkpoint not found at {latest_checkpoint}")
-
-if lora_name:
-    # ... (Код LoRA без изменений, опущен для краткости, если не используется, иначе раскомментируйте оригинальный блок) ...
-    pass 
-
-# Оптимизатор
-if lora_name:
-    trainable_params = [p for p in unet.parameters() if p.requires_grad]
-else:
-    if fbp:
-        trainable_params = list(unet.parameters())
-
-def create_optimizer(name, params):
-    if name == "adam8bit":
-        return bnb.optim.AdamW8bit(
-            params, lr=base_learning_rate, betas=(0.9, betta2), eps=eps, weight_decay=0.01,
-            percentile_clipping=percentile_clipping
-        )
-    elif name == "adam":
-        return torch.optim.AdamW(
-            params, lr=base_learning_rate, betas=(0.9, betta2), eps=1e-8, weight_decay=0.01
-        )
-    elif name == "muon":
-        from muon import MuonWithAuxAdam
-        trainable_params = [p for p in params if p.requires_grad]
-        hidden_weights = [p for p in trainable_params if p.ndim >= 2]
-        hidden_gains_biases = [p for p in trainable_params if p.ndim < 2]
-    
-        param_groups = [ 
-            dict(params=hidden_weights, use_muon=True,
-                lr=1e-3, weight_decay=1e-4),
-            dict(params=hidden_gains_biases, use_muon=False,
-                lr=1e-4, betas=(0.9, 0.95), weight_decay=1e-4),
-        ]
-        optimizer = MuonWithAuxAdam(param_groups)
-        from snooc import SnooC
-        return SnooC(optimizer)
-    else:
-        raise ValueError(f"Unknown optimizer: {name}")
-
-if fbp:
-    optimizer_dict = {p: create_optimizer(optimizer_type, [p]) for p in trainable_params}
-    def optimizer_hook(param):
-        optimizer_dict[param].step()
-        optimizer_dict[param].zero_grad(set_to_none=True)
-    for param in trainable_params:
-        param.register_post_accumulate_grad_hook(optimizer_hook)
-    unet, optimizer = accelerator.prepare(unet, optimizer_dict)
-else:
-    optimizer = create_optimizer(optimizer_type, unet.parameters())
-    def lr_schedule(step):
-        x = step / (total_training_steps * world_size)
-        warmup = warmup_percent
-        if not use_decay:
-            return base_learning_rate
-        if x < warmup:
-            return min_learning_rate + (base_learning_rate - min_learning_rate) * (x / warmup)
-        decay_ratio = (x - warmup) / (1 - warmup)
-        return min_learning_rate + 0.5 * (base_learning_rate - min_learning_rate) * \
-               (1 + math.cos(math.pi * decay_ratio))
-    lr_scheduler = LambdaLR(optimizer, lambda step: lr_schedule(step) / base_learning_rate)
-    unet, optimizer, lr_scheduler = accelerator.prepare(unet, optimizer, lr_scheduler)
-
-if torch_compile:
-    print("compiling")
-    unet = torch.compile(unet)
-    print("compiling - ok")
-
-# Фиксированные семплы
-fixed_samples = get_fixed_samples_by_resolution(dataset)
-
-# --- [UPDATED] Функция для негативного эмбеддинга (возвращает 3 элемента) ---
-def get_negative_embedding(neg_prompt="", batch_size=1):
-    if not neg_prompt:
-        hidden_dim = 2048 
-        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)
-        return empty_emb, empty_mask
-
-    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)
-
-    return uncond_emb, uncond_mask
-    
-# Получаем негативные (пустые) усл��вия для валидации
-uncond_emb, uncond_mask = get_negative_embedding("low quality")
-
-# --- Функция генерации семплов  ---
-@torch.compiler.disable()
-@torch.no_grad()
-def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
-    uncond_emb, uncond_mask = uncond_data
-    
-    original_model = None
-    try:
-        if not torch_compile:
-            original_model = accelerator.unwrap_model(unet, keep_torch_compile=True).eval()
-        else:
-            original_model = unet.eval()
-
-        vae.to(device=device).eval() 
-        
-        all_generated_images = []
-        all_captions = [] 
-        
-        # Распаковываем 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)
-        
-            latents = torch.randn(
-                sample_latents.shape,
-                device=device,
-                dtype=sample_latents.dtype,
-                generator=torch.Generator(device=device).manual_seed(seed)
-            )
-        
-            scheduler.set_timesteps(n_diffusion_steps, device=device)
-        
-            for t in scheduler.timesteps:
-                if guidance_scale != 1:
-                    latent_model_input = torch.cat([latents, latents], dim=0)
-                
-                    # Подготовка батчей для CFG (Negative + Positive)
-                    # 1. Embeddings
-                    curr_batch_size = sample_text_embeddings.shape[0]
-                    seq_len = sample_text_embeddings.shape[1]
-                    hidden_dim = sample_text_embeddings.shape[2]
-                    
-                    neg_emb_batch = uncond_emb[0:1].expand(curr_batch_size, -1, -1)
-                    text_embeddings_batch = torch.cat([neg_emb_batch, sample_text_embeddings], dim=0)
-                    
-                    # 2. Masks
-                    neg_mask_batch = uncond_mask[0:1].expand(curr_batch_size, -1)
-                    attention_mask_batch = torch.cat([neg_mask_batch, sample_mask], dim=0)
-
-                else:
-                    latent_model_input = latents
-                    text_embeddings_batch = sample_text_embeddings
-                    attention_mask_batch = sample_mask
-
-                # Предсказание с передачей всех условий
-                model_out = original_model(
-                    latent_model_input, 
-                    t, 
-                    encoder_hidden_states=text_embeddings_batch, 
-                    encoder_attention_mask=attention_mask_batch,
-                )
-                flow = getattr(model_out, "sample", model_out)
-        
-                if guidance_scale != 1:
-                    flow_uncond, flow_cond = flow.chunk(2)
-                    flow = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
-        
-                latents = scheduler.step(flow, t, latents).prev_sample
-        
-            current_latents = latents
-            if step==0:
-                current_latents = sample_latents
-
-            latent_for_vae = current_latents.detach() / scaling_factor + shift_factor
-            decoded = vae.decode(latent_for_vae.to(torch.float32)).sample
-            decoded_fp32 = decoded.to(torch.float32)
-            
-            for img_idx, img_tensor in enumerate(decoded_fp32):
-                img = (img_tensor / 2 + 0.5).clamp(0, 1).cpu().numpy()
-                img = img.transpose(1, 2, 0)
-    
-                if np.isnan(img).any():
-                    print("NaNs found, saving stopped! Step:", step)
-                pil_img = Image.fromarray((img * 255).astype("uint8"))
-                
-                max_w_overall = max(s[0] for s in fixed_samples_cpu.keys())
-                max_h_overall = max(s[1] for s in fixed_samples_cpu.keys())
-                max_w_overall = max(255, max_w_overall)
-                max_h_overall = max(255, max_h_overall)
-            
-                padded_img = ImageOps.pad(pil_img, (max_w_overall, max_h_overall), color='white')
-                all_generated_images.append(padded_img)
-
-                caption_text = sample_text[img_idx][:300] if img_idx < len(sample_text) else ""
-                all_captions.append(caption_text)
-                
-                sample_path = f"{generated_folder}/{project}_{width}x{height}_{img_idx}.jpg"
-                pil_img.save(sample_path, "JPEG", quality=96)
-        
-        if use_wandb and accelerator.is_main_process:
-            wandb_images = [
-                wandb.Image(img, caption=f"{all_captions[i]}")
-                for i, img in enumerate(all_generated_images)
-            ]
-            wandb.log({"generated_images": wandb_images})
-        if use_comet_ml and accelerator.is_main_process:
-            for i, img in enumerate(all_generated_images):
-                comet_experiment.log_image(
-                    image_data=img,
-                    name=f"step_{step}_img_{i}",
-                    step=step,
-                    metadata={"caption": all_captions[i]}
-                )
-    finally:
-        vae.to("cpu")
-        torch.cuda.empty_cache()
-        gc.collect()
-
-# --------------------------- Генерация сэмплов перед обучением ---------------------------
-if accelerator.is_main_process:
-    if save_model:
-        print("Генерация сэмплов до старта обучения...")
-        generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), 0)
-accelerator.wait_for_everyone()
-
-def save_checkpoint(unet, variant=""):
-    if accelerator.is_main_process:
-        if lora_name:
-            save_lora_checkpoint(unet)
-        else:
-            model_to_save = None
-            if not torch_compile:
-                model_to_save = accelerator.unwrap_model(unet)
-            else:
-                model_to_save = unet
-
-            if variant != "":
-                model_to_save.to(dtype=torch.float16).save_pretrained(
-                    os.path.join(checkpoints_folder, f"{project}"), variant=variant
-                )
-            else:
-                model_to_save.save_pretrained(os.path.join(checkpoints_folder, f"{project}"))
-
-            unet = unet.to(dtype=dtype)
-
-# --------------------------- Тренировочный цикл ---------------------------
-if accelerator.is_main_process:
-    print(f"Total steps per GPU: {total_training_steps}")
-
-epoch_loss_points = []
-progress_bar = tqdm(total=total_training_steps, disable=not accelerator.is_local_main_process, desc="Training", unit="step")
-
-steps_per_epoch = len(dataloader)
-sample_interval = max(1, steps_per_epoch // sample_interval_share)
-min_loss = 4.
-
-for epoch in range(start_epoch, start_epoch + num_epochs):
-    batch_losses = []
-    batch_grads = []
-    batch_sampler.set_epoch(epoch)
-    accelerator.wait_for_everyone()
-    unet.train()
-    
-    for step, (latents, embeddings, attention_mask) in enumerate(dataloader):
-        with accelerator.accumulate(unet):
-            if save_model == False and epoch == 0 and step == 5 :
-                used_gb = torch.cuda.max_memory_allocated() / 1024**3
-                print(f"Шаг {step}: {used_gb:.2f} GB")
-                
-            # шум
-            noise = torch.randn_like(latents, dtype=latents.dtype)
-            # берём t из [0, 1]
-            t = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            #u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            #t = torch.sigmoid(torch.randn_like(u))
-            
-            # интерполяция между x0 и шумом
-            noisy_latents = (1.0 - t.view(-1, 1, 1, 1)) * latents + t.view(-1, 1, 1, 1) * noise
-            # делаем integer timesteps для UNet
-            timesteps = (t * scheduler.config.num_train_timesteps).long()
-            
-            # --- Вызов UNet с маской  ---
-            model_pred = unet(
-                noisy_latents, 
-                timesteps, 
-                encoder_hidden_states=embeddings,
-                encoder_attention_mask=attention_mask
-            ).sample
-            
-            target = noise - latents
-
-            mse_loss = F.mse_loss(model_pred.float(), target.float())
-            mae_loss = F.l1_loss(model_pred.float(), target.float())
-            batch_losses.append(mse_loss.detach().item())
-
-            if (global_step % 100 == 0) or (global_step % sample_interval == 0):
-                accelerator.wait_for_everyone()
-
-            losses_dict = {}
-            losses_dict["mse"] = mse_loss
-            losses_dict["mae"] = mae_loss
-
-            # === Нормализация всех лоссов ===
-            abs_for_norm = {k: losses_dict.get(k, torch.tensor(0.0, device=device)) for k in normalizer.ratios.keys()}
-            total_loss, coeffs, meds = normalizer.update_and_total(abs_for_norm)
-
-            if (global_step % 100 == 0) or (global_step % sample_interval == 0):
-                accelerator.wait_for_everyone()
-                
-            accelerator.backward(total_loss)
-
-            if (global_step % 100 == 0) or (global_step % sample_interval == 0):
-                accelerator.wait_for_everyone()
-                
-            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)
-                    grad = float(grad_val)
-                    optimizer.step()
-                    lr_scheduler.step()
-                    optimizer.zero_grad(set_to_none=True)
-
-            if accelerator.sync_gradients:
-                global_step += 1
-                progress_bar.update(1)
-                if accelerator.is_main_process:
-                    if fbp:
-                        current_lr = base_learning_rate
-                    else:
-                        current_lr = lr_scheduler.get_last_lr()[0]
-                    batch_grads.append(grad)
-    
-                    log_data = {}
-                    log_data["loss_mse"] = mse_loss.detach().item()
-                    log_data["loss_mae"] = mae_loss.detach().item()
-                    log_data["lr"] = current_lr
-                    log_data["grad"] = grad
-                    log_data["loss_norm"] = float(total_loss.item())
-                    for k, c in coeffs.items():
-                        log_data[f"coeff_{k}"] = float(c)
-                    if accelerator.sync_gradients:
-                        if use_wandb:
-                            wandb.log(log_data, step=global_step)
-                        if use_comet_ml:
-                            comet_experiment.log_metrics(log_data, step=global_step)
-
-                    if global_step % sample_interval == 0:
-                        # Передаем tuple (emb, mask) для негатива
-                        if save_model:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
-                        elif epoch % 10 == 0:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask), global_step)
-                        last_n = sample_interval
-                        
-                        if save_model:
-                            has_losses = len(batch_losses) > 0
-                            avg_sample_loss = np.mean(batch_losses[-sample_interval:]) if has_losses else 0.0
-                            last_loss = batch_losses[-1] if has_losses else 0.0
-                            max_loss = max(avg_sample_loss, last_loss)    
-                            should_save = max_loss < min_loss * save_barrier
-                            print(
-                                f"Saving: {should_save} | Max: {max_loss:.4f} | "
-                                f"Last: {last_loss:.4f} | Avg: {avg_sample_loss:.4f}"
-                            )
-                            # 6. Сохранение и обновление
-                            if should_save:
-                                min_loss = max_loss
-                                save_checkpoint(unet)
-
-    if accelerator.is_main_process:
-        avg_epoch_loss = np.mean(batch_losses) if len(batch_losses) > 0 else 0.0
-        avg_epoch_grad = np.mean(batch_grads) if len(batch_grads) > 0 else 0.0
-
-        print(f"\nЭпоха {epoch} завершена. Средний лосс: {avg_epoch_loss:.6f}")
-        log_data_ep = {
-                        "epoch_loss": avg_epoch_loss,
-                        "epoch_grad": avg_epoch_grad,
-                        "epoch": epoch + 1,
-                    }
-        if use_wandb:
-            wandb.log(log_data_ep)
-        if use_comet_ml:
-            comet_experiment.log_metrics(log_data_ep)
-
-if accelerator.is_main_process:
-    print("Обучение завершено! Сохраняем финальную модель...")
-    #if save_model:
-    save_checkpoint(unet,"fp16")
-    if use_comet_ml:
-        comet_experiment.end()
-accelerator.free_memory()
-if torch.distributed.is_initialized():
-    torch.distributed.destroy_process_group()
-    
-print("Готово!")

src/sd15_2048.ipynb

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

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:20f9413ba8b38673409699a49f00a6f9794a0c5153625b422b6151af0df9a940
-size 2047294
+oid sha256:25096d7d566bf784726fbdbd762781a0ff0a25523554272c3c426989633dc969
+size 5145945

train.py

@@ -8,7 +8,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,AutoencoderKLFlux2
+from diffusers import UNet2DConditionModel, AutoencoderKL,AutoencoderKLFlux2,AsymmetricAutoencoderKL,FlowMatchEulerDiscreteScheduler
 from accelerate import Accelerator
 from datasets import load_from_disk
 from tqdm import tqdm
@@ -30,10 +30,10 @@ from transformers import AutoTokenizer, AutoModel
 ds_path = "/workspace/sdxs/datasets/mjnj"
 project = "unet"
 batch_size = 48
-base_learning_rate = 2.7e-5 #2.7e-5
+base_learning_rate = 4e-5 #2.7e-5
 min_learning_rate = 1e-5 #2.7e-5
 num_epochs = 50
-sample_interval_share = 5
+sample_interval_share = 10
 cfg_dropout = 0.15
 max_length = 192
 use_wandb = False
@@ -96,8 +96,8 @@ lora_alpha = 64
 print("init")
 
 loss_ratios = {
-    "mse":   1.25,
-    "mae":   0.25,
+    "mse":   1.5,
+    "mae":   0.5,
 }
 median_coeff_steps = 256 
 
@@ -164,10 +164,12 @@ if accelerator.is_main_process:
 torch.backends.cuda.enable_flash_sdp(True)
 
 # --------------------------- Загрузка моделей ---------------------------
-vae = AutoencoderKL.from_pretrained("vae1x", torch_dtype=dtype).to("cpu").eval()
+#vae = AutoencoderKL.from_pretrained("vae1x", torch_dtype=dtype).to("cpu").eval()
+vae = AsymmetricAutoencoderKL.from_pretrained("vae",torch_dtype=dtype).to(device).eval()
 #vae = AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-dev",subfolder="vae",torch_dtype=dtype).to(device).eval()
 tokenizer = AutoTokenizer.from_pretrained("tokenizer")
 text_model = AutoModel.from_pretrained("text_encoder").to(device).eval()
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("scheduler")
 
 # --- [UPDATED] Функция кодирования текста (с маской и пулингом) ---
 def encode_texts(texts, max_length=max_length):
@@ -237,10 +239,6 @@ if shift_factor is None: shift_factor = 0.0
 scaling_factor = getattr(vae.config, "scaling_factor", 1.0)
 if scaling_factor is None: scaling_factor = 1.0
 
-from diffusers import FlowMatchEulerDiscreteScheduler
-num_train_timesteps = 1000
-scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=num_train_timesteps)
-
 class DistributedResolutionBatchSampler(Sampler):
     def __init__(self, dataset, batch_size, num_replicas, rank, shuffle=True, drop_last=True):
         self.dataset = dataset
@@ -708,7 +706,7 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             if (global_step % 100 == 0) or (global_step % sample_interval == 0):
                 accelerator.wait_for_everyone()
                 
-            accelerator.backward(total_loss)
+            accelerator.backward(mse_loss)
 
             if (global_step % 100 == 0) or (global_step % sample_interval == 0):
                 accelerator.wait_for_everyone()

unet/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dd80ac5e521c295e9b9cc5361127114bcbfea059d76f3e3a80775c91ac666281
-size 1798
+oid sha256:78d4828222ad3a8cddeadf895d9a3afce5c95869d374458dc2c7e5d3b9bf9864
+size 1813

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f3bec6d01b481146de7c8e2adcbce93df1f1cc10ac89834e551fb2d450d286f4
-size 6078588464
+oid sha256:5e903b0e21f57f4ebe996e08e09c76f0605377cde3eea67fe8b9ffce399b153f
+size 3566239360

{sdxs_08b → unet_sdxl5}/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f4c96435f2980db8743704e9361889fb5df8c50443518f76cfe966e8dfc9dc53
-size 1803
+oid sha256:dd80ac5e521c295e9b9cc5361127114bcbfea059d76f3e3a80775c91ac666281
+size 1798

{sdxs_08b → unet_sdxl5}/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9b771189641f108016e1642640e43cb7cc65924a9e6d104cd46831854771cb7b
-size 3376002424
+oid sha256:f3bec6d01b481146de7c8e2adcbce93df1f1cc10ac89834e551fb2d450d286f4
+size 6078588464

vae/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

vae/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e2e5b25e39aec4b6a75e4837adec277dfc830e00992e6ce4dd75eb2627d73197
-size 774
+oid sha256:8bf69a4b1ec4c3b6666326d7d17e98e7f7ed6880084c702101bdb3e75905535c
+size 773

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:31c83db36d96ddfd42003f85abe8c22bd03a07b8174135351345d2726bd75c38
+oid sha256:ffde397a3e78a779adff8ba78297f66d01af5e397512f6ed6d500df30e9833a1
 size 382598708

vae/train_vae_fdl_distil.py

@@ -0,0 +1,651 @@
+# -*- coding: utf-8 -*-
+import os
+import math
+import re
+import torch
+import numpy as np
+import random
+import gc
+from datetime import datetime
+from pathlib import Path
+
+import torchvision.transforms as transforms
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from torch.optim.lr_scheduler import LambdaLR
+from diffusers import AutoencoderKL, AsymmetricAutoencoderKL
+# QWEN: импорт класса
+from diffusers import AutoencoderKLQwenImage
+from diffusers import AutoencoderKLWan
+
+from accelerate import Accelerator
+from PIL import Image, UnidentifiedImageError
+from tqdm import tqdm
+import bitsandbytes as bnb
+import wandb
+import lpips   # pip install lpips
+from FDL_pytorch import FDL_loss # pip install fdl-pytorch
+from collections import deque
+
+# --------------------------- Параметры ---------------------------
+ds_path            = "/workspace/d23"
+project            = "vae5"
+batch_size         = 2
+base_learning_rate = 4e-5
+min_learning_rate  = 2e-5
+num_epochs         = 10
+sample_interval_share = 10
+use_wandb          = True
+save_model         = True
+use_decay          = True
+optimizer_type     = "adam8bit"
+dtype              = torch.float32
+
+model_resolution   = 256
+high_resolution    = 512 
+limit              = 0
+save_barrier       = 1.3
+warmup_percent     = 0.005
+percentile_clipping = 99
+beta2              = 0.997
+eps                = 1e-8
+clip_grad_norm     = 1.0
+mixed_precision    = "no"
+gradient_accumulation_steps = 4
+generated_folder   = "samples"
+save_as            = "vae6"
+num_workers        = 0
+device = None
+
+# --- Режимы обучения ---
+# QWEN: учим только декодер
+train_decoder_only = False
+train_up_only      = False
+full_training      = True  # если True — учим весь VAE и добавляем KL (ниже)
+kl_ratio           = 0.001
+
+# Доли лоссов
+loss_ratios = {
+    "lpips": 0.55,#0.50,
+    "fdl" :  0.05,#0.25,   
+    "edge":  0.05,
+    "mse":   0.10,
+    "mae":   0.04,
+    "kl":    0.001,  # активируем при full_training=True
+    "vae2":  0.199,
+}
+median_coeff_steps = 1000
+
+resize_long_side = 1280  # ресайз длинной стороны исходных картинок
+
+# QWEN: конфиг загрузки модели
+vae_kind      = "kl"  # "qwen" или "kl" (обычный)
+
+Path(generated_folder).mkdir(parents=True, exist_ok=True)
+
+accelerator = Accelerator(
+    mixed_precision=mixed_precision,
+    gradient_accumulation_steps=gradient_accumulation_steps
+)
+device = accelerator.device
+
+# reproducibility
+seed = int(datetime.now().strftime("%Y%m%d"))
+torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
+torch.backends.cudnn.benchmark = False
+
+# --------------------------- WandB ---------------------------
+if use_wandb and accelerator.is_main_process:
+    wandb.init(project=project, config={
+        "batch_size": batch_size,
+        "base_learning_rate": base_learning_rate,
+        "num_epochs": num_epochs,
+        "optimizer_type": optimizer_type,
+        "model_resolution": model_resolution,
+        "high_resolution": high_resolution,
+        "gradient_accumulation_steps": gradient_accumulation_steps,
+        "train_decoder_only": train_decoder_only,
+        "full_training": full_training,
+        "kl_ratio": kl_ratio,
+        "vae_kind": vae_kind,
+    })
+
+# --------------------------- VAE ---------------------------
+def get_core_model(model):
+    m = model
+    # если модель уже обёрнута torch.compile
+    if hasattr(m, "_orig_mod"):
+        m = m._orig_mod
+    return m
+
+def is_video_vae(model) -> bool:
+    # WAN/Qwen — это видео-VAEs
+    if vae_kind in ("wan", "qwen"):
+        return True
+    # fallback по структуре (если понадобится)
+    try:
+        core = get_core_model(model)
+        enc = getattr(core, "encoder", None)
+        conv_in = getattr(enc, "conv_in", None)
+        w = getattr(conv_in, "weight", None)
+        if isinstance(w, torch.nn.Parameter):
+            return w.ndim == 5
+    except Exception:
+        pass
+    return False
+
+# загрузка
+if vae_kind == "qwen":
+    vae = AutoencoderKLQwenImage.from_pretrained("Qwen/Qwen-Image", subfolder="vae")
+else:
+    if vae_kind == "wan":
+        vae = AutoencoderKLWan.from_pretrained(project)
+    else:
+        # старое поведение (пример)
+        if model_resolution==high_resolution:
+            vae = AutoencoderKL.from_pretrained(project)
+        else:
+            vae = AsymmetricAutoencoderKL.from_pretrained(project)
+
+vae = vae.to(dtype)
+
+# --------------------------- VAE2 (Distillation Teacher) ---------------------------
+# Загружаем учителя (SD 1.4) для дистилляции
+print("[INFO] Loading VAE2 (Teacher) for distillation...")
+vae2 = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
+vae2.requires_grad_(False)
+vae2.eval()
+# vae2 перенесем на device позже внутри accelerator.prepare или явно, 
+# но для надежности сделаем это здесь, чтобы не занимать лишнюю память оптимизатором
+vae2.to(device, dtype=dtype)
+
+# Адаптер для проекции 16 каналов студента -> 4 канала учителя
+# Kernel size 1 делает линейную проекцию по пикселям
+distill_adapter = torch.nn.Conv2d(16, 4, kernel_size=1, stride=1, padding=0, bias=True)
+distill_adapter.to(device, dtype=dtype)
+distill_adapter.train() # Адаптер мы обучаем!
+
+
+# torch.compile (опционально)
+if hasattr(torch, "compile"):
+    try:
+        vae = torch.compile(vae)
+    except Exception as e:
+        print(f"[WARN] torch.compile failed: {e}")
+
+# --------------------------- Freeze/Unfreeze ---------------------------
+core = get_core_model(vae)
+
+for p in core.parameters():
+    p.requires_grad = False
+
+unfrozen_param_names = []
+
+if full_training and not train_decoder_only:
+    for name, p in core.named_parameters():
+        p.requires_grad = True
+        unfrozen_param_names.append(name)
+        loss_ratios["kl"] = float(kl_ratio)
+        trainable_module = core
+else:
+    # учим только 0-й блок декодера + post_quant_conv
+    if hasattr(core, "decoder"):
+        if train_up_only:#hasattr(core.decoder, "up_blocks") and len(core.decoder.up_blocks) > 0:
+            # --- только 0-й up_block ---
+            for name, p in core.decoder.up_blocks[0].named_parameters():
+                p.requires_grad = True
+                unfrozen_param_names.append(f"{name}")
+        else:
+            print("Decoder — fallback to full decoder")
+            for name, p in core.decoder.named_parameters():
+                p.requires_grad = True
+                unfrozen_param_names.append(f"decoder.{name}")
+    if hasattr(core, "post_quant_conv"):
+        for name, p in core.post_quant_conv.named_parameters():
+            p.requires_grad = True
+            unfrozen_param_names.append(f"post_quant_conv.{name}")
+    trainable_module = core.decoder if hasattr(core, "decoder") else core
+
+
+print(f"[INFO] Разморожено параметров: {len(unfrozen_param_names)}. Первые 200 имён:")
+for nm in unfrozen_param_names[:200]:
+    print(" ", nm)
+
+# --------------------------- Датасет ---------------------------
+class PngFolderDataset(Dataset):
+    def __init__(self, root_dir, min_exts=('.png',), resolution=1024, limit=0):
+        self.root_dir = root_dir
+        self.resolution = resolution
+        self.paths = []
+        for root, _, files in os.walk(root_dir):
+            for fname in files:
+                if fname.lower().endswith(tuple(ext.lower() for ext in min_exts)):
+                    self.paths.append(os.path.join(root, fname))
+        if limit:
+            self.paths = self.paths[:limit]
+        valid = []
+        for p in self.paths:
+            try:
+                with Image.open(p) as im:
+                    im.verify()
+                valid.append(p)
+            except (OSError, UnidentifiedImageError):
+                continue
+        self.paths = valid
+        if len(self.paths) == 0:
+            raise RuntimeError(f"No valid PNG images found under {root_dir}")
+        random.shuffle(self.paths)
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        p = self.paths[idx % len(self.paths)]
+        with Image.open(p) as img:
+            img = img.convert("RGB")
+            if not resize_long_side or resize_long_side <= 0:
+                return img
+            w, h = img.size
+            long = max(w, h)
+            if long <= resize_long_side:
+                return img
+            scale = resize_long_side / float(long)
+            new_w = int(round(w * scale))
+            new_h = int(round(h * scale))
+            return img.resize((new_w, new_h), Image.BICUBIC)
+
+def random_crop(img, sz):
+    w, h = img.size
+    if w < sz or h < sz:
+        img = img.resize((max(sz, w), max(sz, h)), Image.BICUBIC)
+    x = random.randint(0, max(1, img.width - sz))
+    y = random.randint(0, max(1, img.height - sz))
+    return img.crop((x, y, x + sz, y + sz))
+
+tfm = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+])
+
+dataset = PngFolderDataset(ds_path, min_exts=('.png',), resolution=high_resolution, limit=limit)
+if len(dataset) < batch_size:
+    raise RuntimeError(f"Not enough valid images ({len(dataset)}) to form a batch of size {batch_size}")
+
+def collate_fn(batch):
+    imgs = []
+    for img in batch:
+        img = random_crop(img, high_resolution)
+        imgs.append(tfm(img))
+    return torch.stack(imgs)
+
+dataloader = DataLoader(
+    dataset,
+    batch_size=batch_size,
+    shuffle=True,
+    collate_fn=collate_fn,
+    num_workers=num_workers,
+    pin_memory=True,
+    drop_last=True
+)
+
+# --------------------------- Оптимизатор ---------------------------
+def get_param_groups(module, weight_decay=0.001):
+    no_decay = ["bias", "LayerNorm.weight", "layer_norm.weight", "ln_1.weight", "ln_f.weight"]
+    decay_params, no_decay_params = [], []
+    for n, p in vae.named_parameters():  # глобально по vae, с фильтром requires_grad
+        if not p.requires_grad:
+            continue
+        if any(nd in n for nd in no_decay):
+            no_decay_params.append(p)
+        else:
+            decay_params.append(p)
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+def get_param_groups(module, weight_decay=0.001):
+    no_decay_tokens = ("bias", "norm", "rms", "layernorm")
+    decay_params, no_decay_params = [], []
+    for n, p in module.named_parameters():
+        if not p.requires_grad:
+            continue
+        n_l = n.lower()
+        if any(t in n_l for t in no_decay_tokens):
+            no_decay_params.append(p)
+        else:
+            decay_params.append(p)
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+def create_optimizer(name, param_groups):
+    if name == "adam8bit":
+        return bnb.optim.AdamW8bit(param_groups, lr=base_learning_rate, betas=(0.9, beta2), eps=eps)
+    raise ValueError(name)
+
+param_groups = get_param_groups(get_core_model(vae), weight_decay=0.001)
+
+# --- ИЗМЕНЕНИЕ: Добавляем параметры адаптера в оптимизатор ---
+# Адаптер маленький, weight_decay ему особо не нужен, но пусть будет стандартный
+adapter_params = get_param_groups(distill_adapter, weight_decay=0.001)
+param_groups.extend(adapter_params)
+optimizer = create_optimizer(optimizer_type, param_groups)
+
+# --------------------------- LR schedule ---------------------------
+batches_per_epoch = len(dataloader)
+steps_per_epoch = int(math.ceil(batches_per_epoch / float(gradient_accumulation_steps)))
+total_steps = steps_per_epoch * num_epochs
+
+def lr_lambda(step):
+    if not use_decay:
+        return 1.0
+    x = float(step) / float(max(1, total_steps))
+    warmup = float(warmup_percent)
+    min_ratio = float(min_learning_rate) / float(base_learning_rate)
+    if x < warmup:
+        return min_ratio + (1.0 - min_ratio) * (x / warmup)
+    decay_ratio = (x - warmup) / (1.0 - warmup)
+    return min_ratio + 0.5 * (1.0 - min_ratio) * (1.0 + math.cos(math.pi * decay_ratio))
+
+scheduler = LambdaLR(optimizer, lr_lambda)
+
+# Подготовка
+dataloader, vae, distill_adapter, optimizer, scheduler = accelerator.prepare(
+    dataloader, vae, distill_adapter, optimizer, scheduler
+)
+# vae2 остался обычным тензором на GPU, accelerator его не трогает
+# ----------------------------------------------------------
+
+trainable_params = [p for p in vae.parameters() if p.requires_grad] + \
+                   [p for p in distill_adapter.parameters() if p.requires_grad]
+
+# fdl
+fdl_loss = FDL_loss()
+fdl_loss = fdl_loss.to(accelerator.device)
+
+# --------------------------- LPIPS и вспомогательные ---------------------------
+_lpips_net = None
+def _get_lpips():
+    global _lpips_net
+    if _lpips_net is None:
+        _lpips_net = lpips.LPIPS(net='vgg', verbose=False).eval().to(accelerator.device).eval()
+    return _lpips_net
+
+_sobel_kx = torch.tensor([[[[-1,0,1],[-2,0,2],[-1,0,1]]]], dtype=torch.float32)
+_sobel_ky = torch.tensor([[[[-1,-2,-1],[0,0,0],[1,2,1]]]], dtype=torch.float32)
+def sobel_edges(x: torch.Tensor) -> torch.Tensor:
+    C = x.shape[1]
+    kx = _sobel_kx.to(x.device, x.dtype).repeat(C, 1, 1, 1)
+    ky = _sobel_ky.to(x.device, x.dtype).repeat(C, 1, 1, 1)
+    gx = F.conv2d(x, kx, padding=1, groups=C)
+    gy = F.conv2d(x, ky, padding=1, groups=C)
+    return torch.sqrt(gx * gx + gy * gy + 1e-12)
+
+class MedianLossNormalizer:
+    def __init__(self, desired_ratios: dict, window_steps: int):
+        s = sum(desired_ratios.values())
+        self.ratios = {k: (v / s) if s > 0 else 0.0 for k, v in desired_ratios.items()}
+        self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}
+        self.window = window_steps
+
+    def update_and_total(self, abs_losses: dict):
+        for k, v in abs_losses.items():
+            if k in self.buffers:
+                self.buffers[k].append(float(v.detach().abs().cpu()))
+        meds = {k: (np.median(self.buffers[k]) if len(self.buffers[k]) > 0 else 1.0) for k in self.buffers}
+        coeffs = {k: (self.ratios[k] / max(meds[k], 1e-12)) for k in self.ratios}
+        total = sum(coeffs[k] * abs_losses[k] for k in abs_losses if k in coeffs)
+        return total, coeffs, meds
+
+if full_training and not train_decoder_only:
+    loss_ratios["kl"] = float(kl_ratio)
+normalizer = MedianLossNormalizer(loss_ratios, median_coeff_steps)
+
+# --------------------------- Сэмплы ---------------------------
+@torch.no_grad()
+def get_fixed_samples(n=3):
+    idx = random.sample(range(len(dataset)), min(n, len(dataset)))
+    pil_imgs = [dataset[i] for i in idx]
+    tensors = []
+    for img in pil_imgs:
+        img = random_crop(img, high_resolution)
+        tensors.append(tfm(img))
+    return torch.stack(tensors).to(accelerator.device, dtype)
+
+fixed_samples = get_fixed_samples()
+
+@torch.no_grad()
+def _to_pil_uint8(img_tensor: torch.Tensor) -> Image.Image:
+    arr = ((img_tensor.float().clamp(-1, 1) + 1.0) * 127.5).clamp(0, 255).byte().cpu().numpy().transpose(1, 2, 0)
+    return Image.fromarray(arr)
+
+
+@torch.no_grad()
+def generate_and_save_samples(step=None):
+    try:
+        temp_vae = accelerator.unwrap_model(vae).eval()
+        lpips_net = _get_lpips()
+        with torch.no_grad():
+            orig_high = fixed_samples
+            orig_low = F.interpolate(
+                orig_high,
+                size=(model_resolution, model_resolution),
+                mode="bilinear",
+                align_corners=False
+            )
+            model_dtype = next(temp_vae.parameters()).dtype
+            orig_low = orig_low.to(dtype=model_dtype)
+
+            # Encode/decode с учётом видео-режима
+            if is_video_vae(temp_vae):
+                x_in = orig_low.unsqueeze(2)  # [B,3,1,H,W]
+                enc = temp_vae.encode(x_in)
+                latents_mean = enc.latent_dist.mean
+                dec = temp_vae.decode(latents_mean).sample  # [B,3,1,H,W]
+                rec = dec.squeeze(2)  # [B,3,H,W]
+            else:
+                enc = temp_vae.encode(orig_low)
+                latents_mean = enc.latent_dist.mean
+                rec = temp_vae.decode(latents_mean).sample
+
+        # Подгон размеров, если надо
+        #if rec.shape[-2:] != orig_high.shape[-2:]:
+        #    rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
+
+        # Сохраняем все real/decoded
+        for i in range(rec.shape[0]):
+            real_img = _to_pil_uint8(orig_high[i])
+            dec_img  = _to_pil_uint8(rec[i])
+            real_img.save(f"{generated_folder}/sample_real_{i}.jpg", quality=95)
+            dec_img.save(f"{generated_folder}/sample_decoded_{i}.jpg", quality=95)
+
+        # LPIPS
+        lpips_scores = []
+        for i in range(rec.shape[0]):
+            orig_full = orig_high[i:i+1].to(torch.float32)
+            rec_full  = rec[i:i+1].to(torch.float32)
+            #if rec_full.shape[-2:] != orig_full.shape[-2:]:
+            #    rec_full = F.interpolate(rec_full, size=orig_full.shape[-2:], mode="bilinear", align_corners=False)
+            lpips_val = lpips_net(orig_full, rec_full).item()
+            lpips_scores.append(lpips_val)
+        avg_lpips = float(np.mean(lpips_scores))
+
+        # W&B логирование
+        if use_wandb and accelerator.is_main_process:
+            log_data = {"lpips_mean": avg_lpips}
+            for i in range(rec.shape[0]):
+                log_data[f"sample/real_{i}"] = wandb.Image(f"{generated_folder}/sample_real_{i}.jpg", caption=f"real_{i}")
+                log_data[f"sample/decoded_{i}"] = wandb.Image(f"{generated_folder}/sample_decoded_{i}.jpg", caption=f"decoded_{i}")
+            wandb.log(log_data, step=step)
+
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+if accelerator.is_main_process and save_model:
+    print("Генерация сэмплов до старта обучения...")
+    generate_and_save_samples(0)
+
+accelerator.wait_for_everyone()
+
+# --------------------------- Тренировка ---------------------------
+progress = tqdm(total=total_steps, disable=not accelerator.is_local_main_process)
+global_step = 0
+min_loss = float("inf")
+sample_interval = max(1, total_steps // max(1, sample_interval_share * num_epochs))
+
+for epoch in range(num_epochs):
+    vae.train()
+    batch_losses, batch_grads = [], []
+    track_losses = {k: [] for k in loss_ratios.keys()}
+
+    for imgs in dataloader:
+        with accelerator.accumulate(vae):
+            imgs = imgs.to(accelerator.device)
+
+            if high_resolution != model_resolution:
+                imgs_low = F.interpolate(imgs, size=(model_resolution, model_resolution), mode="bilinear", align_corners=False)
+            else:
+                imgs_low = imgs
+
+            model_dtype = next(vae.parameters()).dtype
+            imgs_low_model = imgs_low.to(dtype=model_dtype) if imgs_low.dtype != model_dtype else imgs_low
+
+            # QWEN: encode/decode с T=1
+            if is_video_vae(vae):
+                x_in = imgs_low_model.unsqueeze(2)             # [B,3,1,H,W]
+                enc = vae.encode(x_in)
+                latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
+                dec = vae.decode(latents).sample               # [B,3,1,H,W]
+                rec = dec.squeeze(2)                           # [B,3,H,W]
+            else:
+                enc = vae.encode(imgs_low_model)
+                latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
+                rec = vae.decode(latents).sample
+
+            #if rec.shape[-2:] != imgs.shape[-2:]:
+            #    rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
+
+            rec_f32 = rec.to(torch.float32)
+            imgs_f32 = imgs.to(torch.float32)
+
+            # ... (предыдущий код получения rec, rec_f32, imgs_f32) ...
+
+            # --- ИЗМЕНЕНИЕ: Расчет VAE2 Distillation Loss ---
+            # 1. Получаем латенты учителя (SD). 
+            #    SD VAE ожидает вход [-1, 1], у нас imgs_low_model уже нормализован так же.
+            with torch.no_grad():
+                # scale_factor=0.18215 в SD обычно применяется ПОСЛЕ энкодера для диффузии.
+                # Но для дистилляции мы можем сравнивать "сырые" распределения (moments).
+                # Главное сравнивать яблоки с яблоками. Берем .mean (детерминированный выход).
+                teacher_dist = vae2.encode(imgs_low_model).latent_dist
+                teacher_mean = teacher_dist.mean
+
+            # 2. Получаем латенты студента. 
+            #    Они уже посчитаны выше в переменной 'enc'
+            #    enc.latent_dist.mean - это 16 каналов
+            student_mean = enc.latent_dist.mean
+
+            # 3. Проецируем 16 -> 4 через наш обучаемый адаптер
+            #    distill_adapter уже обернут акселератором, dtype будет правильный
+            student_projected = distill_adapter(student_mean)
+
+            # 4. Считаем MSE между проекцией студента и учителем
+            #    Оба тензора должны быть [B, 4, H_lat, W_lat]
+            loss_distill = F.mse_loss(student_projected.float(), teacher_mean.float())
+            # ------------------------------------------------
+
+            abs_losses = {
+                "mae":   F.l1_loss(rec_f32, imgs_f32),
+                "mse":   F.mse_loss(rec_f32, imgs_f32),
+                "lpips": _get_lpips()(rec_f32, imgs_f32).mean(),
+                "fdl":   fdl_loss(rec_f32, imgs_f32),
+                "edge":  F.l1_loss(sobel_edges(rec_f32), sobel_edges(imgs_f32)),
+                "vae2":  loss_distill,  # <--- Добавляем в словарь лоссов
+            }
+
+            if full_training and not train_decoder_only:
+                mean   = enc.latent_dist.mean
+                logvar = enc.latent_dist.logvar
+                kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
+                abs_losses["kl"] = kl
+            else:
+                abs_losses["kl"] = torch.tensor(0.0, device=accelerator.device, dtype=torch.float32)
+
+            total_loss, coeffs, meds = normalizer.update_and_total(abs_losses)
+
+            if torch.isnan(total_loss) or torch.isinf(total_loss):
+                raise RuntimeError("NaN/Inf loss")
+
+            accelerator.backward(total_loss)
+
+            grad_norm = torch.tensor(0.0, device=accelerator.device)
+            if accelerator.sync_gradients:
+                grad_norm = accelerator.clip_grad_norm_(trainable_params, clip_grad_norm)
+                optimizer.step()
+                scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+                global_step += 1
+                progress.update(1)
+
+            if accelerator.is_main_process:
+                try:
+                    current_lr = optimizer.param_groups[0]["lr"]
+                except Exception:
+                    current_lr = scheduler.get_last_lr()[0]
+
+                batch_losses.append(total_loss.detach().item())
+                batch_grads.append(float(grad_norm.detach().cpu().item()) if isinstance(grad_norm, torch.Tensor) else float(grad_norm))
+                for k, v in abs_losses.items():
+                    track_losses[k].append(float(v.detach().item()))
+
+                if use_wandb and accelerator.sync_gradients:
+                    log_dict = {
+                        "total_loss": float(total_loss.detach().item()),
+                        "learning_rate": current_lr,
+                        "epoch": epoch,
+                        "grad_norm": batch_grads[-1],
+                    }
+                    for k, v in abs_losses.items():
+                        log_dict[f"loss_{k}"] = float(v.detach().item())
+                    for k in coeffs:
+                        log_dict[f"coeff_{k}"] = float(coeffs[k])
+                        log_dict[f"median_{k}"] = float(meds[k])
+                    wandb.log(log_dict, step=global_step)
+
+            if global_step > 0 and global_step % sample_interval == 0:
+                if accelerator.is_main_process:
+                    generate_and_save_samples(global_step)
+                accelerator.wait_for_everyone()
+
+                n_micro = sample_interval * gradient_accumulation_steps
+                avg_loss = float(np.mean(batch_losses[-n_micro:])) if len(batch_losses) >= n_micro else float(np.mean(batch_losses)) if batch_losses else float("nan")
+                avg_grad = float(np.mean(batch_grads[-n_micro:])) if len(batch_grads) >= 1 else float(np.mean(batch_grads)) if batch_grads else 0.0
+
+                if accelerator.is_main_process:
+                    print(f"Epoch {epoch} step {global_step} loss: {avg_loss:.6f}, grad_norm: {avg_grad:.6f}, lr: {current_lr:.9f}")
+                    if save_model and avg_loss < min_loss * save_barrier:
+                        min_loss = avg_loss
+                        accelerator.unwrap_model(vae).save_pretrained(save_as)
+                    if use_wandb:
+                        wandb.log({"interm_loss": avg_loss, "interm_grad": avg_grad}, step=global_step)
+
+    if accelerator.is_main_process:
+        epoch_avg = float(np.mean(batch_losses)) if batch_losses else float("nan")
+        print(f"Epoch {epoch} done, avg loss {epoch_avg:.6f}")
+        if use_wandb:
+            wandb.log({"epoch_loss": epoch_avg, "epoch": epoch + 1}, step=global_step)
+
+# --------------------------- Финальное сохранение ---------------------------
+if accelerator.is_main_process:
+    print("Training finished – saving final model")
+    if save_model:
+        accelerator.unwrap_model(vae).save_pretrained(save_as)
+
+accelerator.free_memory()
+if torch.distributed.is_initialized():
+    torch.distributed.destroy_process_group()
+print("Готово!")

vae2x/config.json

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

vae2x/diffusion_pytorch_model.safetensors

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