2602

dataset.py

@@ -18,12 +18,12 @@ 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
-min_size = 320 #384 #320 #192 #256 #192
-max_size = 640 #768 #640 #384 #256 #384
-step = 32 #64
+min_size = 640 #320 #384 #320 #192 #256 #192
+max_size = 1280 #640 #768 #640 #384 #256 #384
+step = 64
 empty_share = 0.0
 limit = 0
 # Основная процедура обработки
@@ -43,20 +43,9 @@ def clear_cuda_memory():
 def load_models():
     print("Загрузка моделей...")
     #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)
-    #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 = AutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
+    return vae
+
 vae = load_models()
 
 shift_factor = getattr(vae.config, "shift_factor", 0.0)
@@ -67,8 +56,11 @@ 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)
+mean = getattr(vae.config, "latents_mean", None)
+std = getattr(vae.config, "latents_std", None)
+if mean is not None and std is not None:
+    latents_std = torch.tensor(std, device=device, dtype=dtype).view(1, len(std), 1, 1)
+    latents_mean = torch.tensor(mean, device=device, dtype=dtype).view(1, len(mean), 1, 1)
     
 # ---------------- 3️⃣ Трансформации ----------------
 def get_image_transform(min_size=256, max_size=512, step=64):
@@ -126,50 +118,6 @@ def get_image_transform(min_size=256, max_size=512, step=64):
     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("."):
@@ -200,42 +148,6 @@ 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)
@@ -269,20 +181,19 @@ 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)
+            if latents_mean is not None and latents_std is not None:
+                posteriors = (posteriors - latents_mean) / latents_std
+            posteriors = (posteriors - shift_factor) / scaling_factor
         
-        latents_np = image_latents.to(dtype).cpu().numpy()
+        latents_np = posteriors.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

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:491d2fb7582ad01f0c930c28f58476cf8f46e0900417c56ce3a06e8f63eb19df
+oid sha256:12e68807a96fc1f2dd993c400888104240383fd74379a2f083952996c0d60c13
 size 5946605448