flux

dataset-Copy1.py

@@ -0,0 +1,412 @@
+# pip install flash-attn --no-build-isolation
+import torch
+import os
+import gc
+import numpy as np
+import random
+import json
+import shutil
+import time
+
+from datasets import Dataset, load_from_disk, concatenate_datasets
+from diffusers import AutoencoderKL,AutoencoderKLWan,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
+
+# ---------------- 1️⃣ Настройки ----------------
+dtype = torch.float32
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+batch_size = 5
+min_size = 320 #384 #320 #192 #256 #192
+max_size = 640 #768 #640 #384 #256 #384
+step = 32 #64
+empty_share = 0.0
+limit = 0
+# Основная процедура обработки
+folder_path = "/workspace/dataset" #alchemist"
+save_path = "/workspace/dataset_640" #"alchemist"
+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 = AsymmetricAutoencoderKL.from_pretrained("AiArtLab/sdxs-1b",subfolder="vae",torch_dtype=dtype).to(device).eval()
+
+    #model_name = "Qwen/Qwen3-0.6B"
+    #tokenizer = AutoTokenizer.from_pretrained(model_name)
+    #model = AutoModelForCausalLM.from_pretrained(
+    #    model_name,
+    #    torch_dtype=dtype,
+    #    device_map=device
+    #).eval()
+    #tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B', padding_side='left')
+    #model = AutoModel.from_pretrained('Qwen/Qwen3-Embedding-0.6B').to("cuda")
+    return vae#, model, tokenizer
+
+#vae, model, tokenizer = load_models()
+vae = load_models()
+
+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
+    
+latents_mean = getattr(vae.config, "latents_mean", None)
+latents_std = getattr(vae.config, "latents_std", None)
+    
+# ---------------- 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 last_token_pool(last_hidden_states: torch.Tensor,
+                    attention_mask: torch.Tensor) -> torch.Tensor:
+    # Определяем, есть ли left padding
+    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+    if left_padding:
+        return last_hidden_states[:, -1]
+    else:
+        sequence_lengths = attention_mask.sum(dim=1) - 1
+        batch_size = last_hidden_states.shape[0]
+        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
+
+def encode_texts_batch(texts, tokenizer, model, device="cuda", max_length=150, normalize=False):
+    with torch.inference_mode():
+        # Токенизация
+        batch = tokenizer(
+            texts,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=max_length
+        ).to(device)
+
+        # Прогон через модель
+        #outputs = model(**batch)
+
+        # Пулинг по last token
+        #embeddings = last_token_pool(outputs.last_hidden_state, batch["attention_mask"])
+
+        # L2-нормализация (опционально, обычно нужна для семантического поиска)
+        #if normalize:
+        #    embeddings = F.normalize(embeddings, p=2, dim=1)
+
+        # Прогон через базовую модель (внутри CausalLM)
+        outputs = model.model(**batch, output_hidden_states=True)
+
+        # Берем последний слой (эмбеддинги всех токенов)
+        hidden_states = outputs.hidden_states[-1]  # [B, L, D]
+
+        # Можно применить нормализацию по каждому токену (как в CLIP)
+        if normalize:
+            hidden_states = F.normalize(hidden_states, p=2, dim=-1)
+
+    return hidden_states.cpu().numpy() # embeddings.unsqueeze(1).cpu().numpy()
+
+def clean_label(label):
+    label = label.replace("Image 1", "").replace("Image 2", "").replace("Image 3", "").replace("Image 4", "").replace("The image depicts ","").replace("The image presents ","").replace("The image features ","").replace("The image portrays ","").replace("The image is ","").strip()
+    if label.startswith("."):
+        label = label[1:].lstrip()
+    return label
+
+def process_labels_for_guidance(original_labels, prob_to_make_empty=0.01):
+    """
+    Обрабатывает список меток для classifier-free guidance.
+
+    С вероятностью prob_to_make_empty:
+    - Метка в первом списке заменяется на пустую строку.
+    - К метке во втором списке добавляется префикс "zero:".
+
+    В противном случае метки в обоих списках остаются оригинальными.
+
+    """
+    labels_for_model = []
+    labels_for_logging = []
+
+    for label in original_labels:
+        if random.random() < prob_to_make_empty:
+            labels_for_model.append("")  # Заменяем на пустую строку для модели
+            labels_for_logging.append(f"zero: {label}") # Добавляем префикс для логгирования
+        else:
+            labels_for_model.append(label) # Оставляем оригинальную метку для модели
+            labels_for_logging.append(label) # Оставляем оригинальную метку для логгирования
+
+    return labels_for_model, labels_for_logging
+
+def encode_to_latents(images, texts):
+    transform = get_image_transform(min_size, max_size, step)
+    
+    try:
+        # Обработка изображений (все одинакового размера)
+        transformed_tensors = []
+        pil_images = []
+        widths, heights = [], []
+        
+        # Применяем трансформацию ко всем изображениям
+        for img in images:
+            try:
+                t_img, pil_img, w, h = transform(img)
+                transformed_tensors.append(t_img)
+                pil_images.append(pil_img)
+                widths.append(w)
+                heights.append(h)
+            except Exception as e:
+                print(f"Ошибка трансформации: {e}")
+                continue
+
+        if not transformed_tensors:
+            return None
+
+        # Создаём батч
+        batch_tensor = torch.stack(transformed_tensors).to(device, dtype)
+        if batch_tensor.ndim==5:
+            batch_tensor = batch_tensor.unsqueeze(2)  # [B, C, 1, H, W]
+        
+        # Кодируем батч
+        with torch.no_grad():
+            posteriors = vae.encode(batch_tensor).latent_dist.mode()
+            latents = (posteriors - shift_factor) / scaling_factor
+        
+        latents_np = latents.to(dtype).cpu().numpy()
+
+        # Обрабатываем тексты
+        text_labels = [clean_label(text) for text in texts]
+
+        model_prompts, text_labels = process_labels_for_guidance(text_labels, empty_share)                
+        #embeddings = encode_texts_batch(model_prompts, tokenizer, model)
+
+        return {
+            "vae": latents_np,
+            #"embeddings": embeddings,
+            "text": text_labels,
+            "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=10000, 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=20000, batch_size=batch_size)
+
+# Удаление  папки
+try:
+    shutil.rmtree(folder_path)
+    print(f"✅ Папка {folder_path} успешно удалена")
+except Exception as e:
+    print(f"⚠️ Ошибка при удалении папки: {e}")
+    
+# Объединение чанков в финальный датасет
+combine_chunks(temp_path, save_path)
+
+# Удаление временной папки
+try:
+    shutil.rmtree(temp_path)
+    print(f"✅ Временная папка {temp_path} успешно удалена")
+except Exception as e:
+    print(f"⚠️ Ошибка при удалении временной папки: {e}")

dataset.py

@@ -9,7 +9,7 @@ import shutil
 import time
 
 from datasets import Dataset, load_from_disk, concatenate_datasets
-from diffusers import AutoencoderKL,AutoencoderKLWan,AsymmetricAutoencoderKL
+from diffusers import AutoencoderKL,AutoencoderKLWan,AsymmetricAutoencoderKL,AutoencoderKLFlux2
 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
@@ -18,17 +18,17 @@ from tqdm import tqdm
 from datetime import timedelta
 
 # ---------------- 1️⃣ Настройки ----------------
-dtype = torch.float32
+dtype = torch.float16
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-batch_size = 5
+batch_size = 1
 min_size = 320 #384 #320 #192 #256 #192
 max_size = 640 #768 #640 #384 #256 #384
 step = 32 #64
 empty_share = 0.0
 limit = 0
 # Основная процедура обработки
-folder_path = "/workspace/dataset" #alchemist"
-save_path = "/workspace/dataset_640" #"alchemist"
+folder_path = "/home/recoilme/dataset/butterflies/butterfly" #alchemist"
+save_path = "/home/recoilme/sdxs-1b/datasets/butterfly_flux32" #"alchemist"
 os.makedirs(save_path, exist_ok=True)
 
 # Функция для очистки CUDA памяти
@@ -42,7 +42,8 @@ def clear_cuda_memory():
 # ---------------- 2️⃣ Загрузка моделей ----------------
 def load_models():
     print("Загрузка моделей...")
-    vae = AsymmetricAutoencoderKL.from_pretrained("AiArtLab/sdxs-1b",subfolder="vae",torch_dtype=dtype).to(device).eval()
+    #vae = AsymmetricAutoencoderKL.from_pretrained("AiArtLab/sdxs-1b",subfolder="vae",torch_dtype=dtype).to(device).eval()
+    vae = AutoencoderKLFlux2.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
 
     #model_name = "Qwen/Qwen3-0.6B"
     #tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -199,6 +200,43 @@ def process_labels_for_guidance(original_labels, prob_to_make_empty=0.01):
 
     return labels_for_model, labels_for_logging
 
+def _patchify_latents(latents):
+    batch_size, num_channels_latents, height, width = latents.shape
+    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 1, 3, 5, 2, 4)
+    latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2)
+    return latents
+
+@staticmethod
+def _unpatchify_latents(latents):
+    batch_size, num_channels_latents, height, width = latents.shape
+    latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
+    latents = latents.permute(0, 1, 4, 2, 5, 3)
+    latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
+    return latents
+
+def flux_encode(vae,latents):
+    # patch
+    image_latents = _patchify_latents(latents)
+    # norm
+    latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+    latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps)
+    latents = (image_latents - latents_bn_mean) / latents_bn_std
+    # unpatch
+    latents = _unpatchify_latents(latents)
+    return latents
+    
+def flux_decode(vae,latents):
+    # patch
+    image_latents = _patchify_latents(latents)
+    # norm
+    latents_bn_mean = vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+    latents_bn_std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps)
+    latents = image_latents * latents_bn_std + latents_bn_mean
+    # unpatch
+    latents = _unpatchify_latents(latents)
+    return latents
+
 def encode_to_latents(images, texts):
     transform = get_image_transform(min_size, max_size, step)
     
@@ -232,6 +270,7 @@ def encode_to_latents(images, texts):
         with torch.no_grad():
             posteriors = vae.encode(batch_tensor).latent_dist.mode()
             latents = (posteriors - shift_factor) / scaling_factor
+            image_latents = flux_encode(vae, latents)
         
         latents_np = latents.to(dtype).cpu().numpy()
 

samples/unet_320x640_0.jpg → model_index-Copy1.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8f1c5f47692d66d5d06f0324adba126a88fb101ee8bc41400135d52b5afa3645
-size 95393
+oid sha256:e7262c65619c10d3d77607525c101c525b6b6a9a3af89f503f35b42e91dd88e2
+size 417

model_index.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e7262c65619c10d3d77607525c101c525b6b6a9a3af89f503f35b42e91dd88e2
-size 417
+oid sha256:b6d71e1f562601e1a6bfcd1f0f7f81e021003b5512637d1c14dd77aba88144c8
+size 412

pipeline_sdxs-Copy1.py

@@ -0,0 +1,213 @@
+import torch
+import numpy as np
+from PIL import Image
+from typing import List, Union, Optional, Tuple
+from dataclasses import dataclass
+
+from diffusers import DiffusionPipeline
+from diffusers.utils import BaseOutput
+from tqdm import tqdm
+
+@dataclass
+class SdxsPipelineOutput(BaseOutput):
+    images: Union[List[Image.Image], np.ndarray]
+
+class SdxsPipeline(DiffusionPipeline):
+    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler):
+        super().__init__()
+        self.register_modules(
+            vae=vae, 
+            text_encoder=text_encoder, 
+            tokenizer=tokenizer, 
+            unet=unet, 
+            scheduler=scheduler
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        
+    def preprocess_image(self, image: Image.Image, width: int, height: int):
+        """Ресайз и центрированный кроп изображения для асимметричного VAE."""
+        # Для энкодера с масштабом 8
+        target_height = ((height // self.vae_scale_factor) * self.vae_scale_factor)//2
+        target_width = ((width // self.vae_scale_factor) * self.vae_scale_factor)//2
+        
+        w, h = image.size
+        aspect_ratio = target_width / target_height
+        
+        if w / h > aspect_ratio:
+            new_w = int(h * aspect_ratio)
+            left = (w - new_w) // 2
+            image = image.crop((left, 0, left + new_w, h))
+        else:
+            new_h = int(w / aspect_ratio)
+            top = (h - new_h) // 2
+            image = image.crop((0, top, w, top + new_h))
+        
+        image = image.resize((target_width, target_height), resample=Image.LANCZOS)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[None].transpose(0, 3, 1, 2)  # [1, C, H, W]
+        image = torch.from_numpy(image)
+        return 2.0 * image - 1.0  # [-1, 1]
+
+    def encode_prompt(self, prompt, negative_prompt, device, dtype):
+        def get_single_encode(texts, is_negative=False):
+            if texts is None or texts == "":
+                hidden_dim = self.text_encoder.config.hidden_size
+                shape = (1, self.text_encoder.config.max_position_embeddings, hidden_dim)
+                emb = torch.zeros(shape, dtype=dtype, device=device)
+                mask = torch.ones((1, self.text_encoder.config.max_position_embeddings), dtype=torch.int64, device=device)
+                return emb, mask
+
+            if isinstance(texts, str):
+                texts = [texts]
+
+            with torch.no_grad():
+                toks = self.tokenizer(
+                    texts,
+                    padding="max_length",
+                    max_length=self.text_encoder.config.max_position_embeddings,
+                    truncation=True,
+                    return_tensors="pt"
+                ).to(device)
+
+                outputs = self.text_encoder(
+                    input_ids=toks.input_ids, 
+                    attention_mask=toks.attention_mask,
+                    output_hidden_states=True
+                )
+
+                layer_index = -2
+                prompt_embeds = outputs.hidden_states[layer_index]
+                final_layer_norm = self.text_encoder.text_model.final_layer_norm
+                prompt_embeds = final_layer_norm(prompt_embeds)
+
+                return prompt_embeds, toks.attention_mask
+
+        pos_embeds, pos_mask = get_single_encode(prompt)
+        neg_embeds, neg_mask = get_single_encode(negative_prompt, is_negative=True)
+
+        batch_size = pos_embeds.shape[0]
+        if neg_embeds.shape[0] != batch_size:
+            neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
+            neg_mask = neg_mask.repeat(batch_size, 1)
+
+        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
+        final_mask = torch.cat([neg_mask, pos_mask], dim=0)
+
+        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        image: Optional[Union[Image.Image, List[Image.Image]]] = None,
+        coef: float = 0.97,                    # ← strength (0.0 = оригинал, 1.0 = полный шум)
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 4.0,
+        generator: Optional[torch.Generator] = None,
+        seed: Optional[int] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        # structure_preservation оставляем для совместимости, но теперь он почти не нужен
+        structure_preservation: float = 0.0,   # 0.0 = стандартный линейный путь (лучше всего)
+        **kwargs,
+    ):
+        device = self.device
+        dtype = self.unet.dtype
+    
+        if generator is None and seed is not None:
+            generator = torch.Generator(device=device).manual_seed(seed)
+    
+        # 1. Encode prompt (твой код оставляем без изменений)
+        text_embeddings, attention_mask = self.encode_prompt(
+            prompt, negative_prompt, device, dtype
+        )
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    
+        # 2. Scheduler timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+    
+        # ==================== IMG2IMG БЛОК (НОВАЯ ВЕРСИЯ) ====================
+        if image is not None:
+            # --- Подготовка изображения ---
+            if isinstance(image, Image.Image):
+                image_tensor = self.preprocess_image(image, width, height).to(device, self.vae.dtype)
+            else:
+                image_tensor = self.preprocess_image(image[0], width, height).to(device, self.vae.dtype)
+    
+            # --- Кодируем в latent ---
+            latents_clean = self.vae.encode(image_tensor).latent_dist.sample(generator=generator)
+            vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+            vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
+            latents_clean = (latents_clean - vae_shift_factor) / vae_scaling_factor
+            latents_clean = latents_clean.to(dtype)
+    
+            # --- Добавляем шум по Rectified Flow формуле ---
+            noise = torch.randn_like(latents_clean)
+    
+            # coef = strength (0.0 → оригинал, 1.0 → чистый шум)
+            sigma = coef                                      # в Flow Matching sigma = t
+            if hasattr(self.scheduler, "sigma_shift"):        # если есть shift (Flux-style)
+                sigma = self.scheduler.sigma_shift(sigma)
+    
+            latents = (1.0 - sigma) * latents_clean + sigma * noise
+    
+            # Обрезаем timesteps начиная с текущего sigma
+            init_timestep = int(num_inference_steps * coef)
+            t_start = max(num_inference_steps - init_timestep, 0)
+            timesteps = timesteps[t_start:]
+    
+            #print(f"img2img → strength={coef:.2f}, sigma={sigma:.3f}, steps={len(timesteps)}")
+    
+        else:
+            # txt2img — оставляем как было
+            vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+            vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
+            latent_h = height // self.vae_scale_factor
+            latent_w = width // self.vae_scale_factor
+    
+            latents = torch.randn(
+                (batch_size, self.unet.config.in_channels, latent_h, latent_w),
+                generator=generator, device=device, dtype=dtype
+            )
+    
+        # ==================== DENOISING LOOP (одинаковый для txt2img и img2img) ====================
+        for i, t in enumerate(tqdm(timesteps, desc="Sampling")):
+            latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents
+    
+            model_out = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                encoder_attention_mask=attention_mask,
+                return_dict=False,
+            )[0]
+    
+            if guidance_scale > 1.0:
+                flow_uncond, flow_cond = model_out.chunk(2)
+                model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
+    
+            # Важно: используем scheduler.step — он сам знает, что делать с velocity
+            latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]
+    
+        # ==================== DECODE ====================
+        if output_type == "latent":
+            return SdxsPipelineOutput(images=latents)
+    
+        latents = latents * vae_scaling_factor + vae_shift_factor
+        image_output = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
+    
+        image_output = (image_output.clamp(-1, 1) + 1) / 2
+        image_np = image_output.cpu().permute(0, 2, 3, 1).float().numpy()
+    
+        if output_type == "pil":
+            images = [(Image.fromarray((img * 255).round().astype("uint8"))) for img in image_np]
+        else:
+            images = image_np
+    
+        if not return_dict:
+            return images
+        return SdxsPipelineOutput(images=images)

pipeline_sdxs.py

@@ -22,7 +22,7 @@ class SdxsPipeline(DiffusionPipeline):
             unet=unet, 
             scheduler=scheduler
         )
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.vae_scale_factor = 16 #2 ** (len(self.vae.config.block_out_channels) - 1)
         
     def preprocess_image(self, image: Image.Image, width: int, height: int):
         """Ресайз и центрированный кроп изображения для асимметричного VAE."""
@@ -48,6 +48,55 @@ class SdxsPipeline(DiffusionPipeline):
         image = torch.from_numpy(image)
         return 2.0 * image - 1.0  # [-1, 1]
 
+    @staticmethod
+    def _patchify_latents(latents):
+        batch_size, num_channels_latents, height, width = latents.shape
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 1, 3, 5, 2, 4)
+        latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2)
+        return latents
+    
+    @staticmethod
+    def _unpatchify_latents(latents):
+        batch_size, num_channels_latents, height, width = latents.shape
+        latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2)
+        return latents
+
+    def flux_encode(self, latents):
+        # 1. Patchify
+        image_latents = self._patchify_latents(latents)
+        
+        # 2. Normalization
+        # Достаем параметры из self.vae
+        bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+        bn_var = self.vae.bn.running_var.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+        eps = getattr(self.vae.config, "batch_norm_eps", 1e-5)
+        
+        latents_bn_std = torch.sqrt(bn_var + eps)
+        latents = (image_latents - bn_mean) / latents_bn_std
+        
+        # 3. Unpatchify
+        latents = self._unpatchify_latents(latents)
+        return latents
+
+    def flux_decode(self, latents):
+        # 1. Patchify
+        image_latents = self._patchify_latents(latents)
+        
+        # 2. De-normalization
+        bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+        bn_var = self.vae.bn.running_var.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype)
+        eps = getattr(self.vae.config, "batch_norm_eps", 1e-5)
+        
+        latents_bn_std = torch.sqrt(bn_var + eps)
+        latents = image_latents * latents_bn_std + bn_mean
+        
+        # 3. Unpatchify
+        latents = self._unpatchify_latents(latents)
+        return latents
+
     def encode_prompt(self, prompt, negative_prompt, device, dtype):
         def get_single_encode(texts, is_negative=False):
             if texts is None or texts == "":
@@ -198,6 +247,15 @@ class SdxsPipeline(DiffusionPipeline):
             return SdxsPipelineOutput(images=latents)
     
         latents = latents * vae_scaling_factor + vae_shift_factor
+        latents = self.flux_decode(latents)
+        #latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+        #latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
+        #    latents.device, latents.dtype
+        #)
+        #latents = latents * latents_bn_std + latents_bn_mean
+        #latents = self._unpatchify_latents(latents)
+
+        image = self.vae.decode(latents, return_dict=False)[0]
         image_output = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
     
         image_output = (image_output.clamp(-1, 1) + 1) / 2

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:99418c2b7f273b90540fbefb47fd60793d6973d4d8238c78ceba86716c36a5d6
-size 5313615
+oid sha256:0a1db3927bbb12c88ac07d9bf9b8006cc095417b561240413be0c9956618046d
+size 2126983

train_flux.py

@@ -33,10 +33,10 @@ from transformers import AutoTokenizer, AutoModel
 ds_path = "/workspace/sdxs-1b/datasets/mjnj_640_flux2"
 project = "unet"
 ## total batch (split // num `GPU)
-batch_size = 12
+batch_size = 42
 base_learning_rate = 3e-5
 min_learning_rate = 1e-5
-num_epochs = 1
+num_epochs = 8
 sample_interval_share = 20
 cfg_dropout = 0.10
 max_length = 248
@@ -761,4 +761,4 @@ accelerator.free_memory()
 if torch.distributed.is_initialized():
     torch.distributed.destroy_process_group()
     
-print("Готово!")
+print("Готово!")

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d1a8c9098a3ea786f6c555bfae1d106df7c128e5702314e28c5053938499e724
+oid sha256:a0bd7c0c156cd06520c8925860a11bf0c7796bc2bf22f5febeb6696dc82896fd
 size 5935560296