2603

pipeline_sdxs-Copy1.py

@@ -1,369 +0,0 @@
-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
-from transformers import Qwen3ForCausalLM, Qwen2Tokenizer
-
-@dataclass
-class SdxsPipelineOutput(BaseOutput):
-    images: Union[List[Image.Image], np.ndarray]
-    prompt: Optional[Union[str, List[str]]] = None  # Возврат улучшенного промпта
-
-class SdxsPipeline(DiffusionPipeline):
-    def __init__(self, vae, text_encoder, text_encoder2, tokenizer, tokenizer2, unet, scheduler):
-        super().__init__()
-        self.register_modules(
-            vae=vae, 
-            text_encoder=text_encoder,
-            text_encoder2=text_encoder2, 
-            tokenizer=tokenizer,
-            tokenizer2=tokenizer2,
-            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)
-        target_width = ((width // self.vae_scale_factor) * self.vae_scale_factor)
-        
-        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]
-
-    @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
-        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):
-            if not texts:
-                texts = [""]
-            elif 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
-                )
-
-                # 1. Берем -2 слой [Batch, Seq, Dim]
-                hidden = outputs.hidden_states[-2]
-
-                # 2. Достаем pooled вектор (последний токен) [Batch, Dim]
-                seq_lens = toks.attention_mask.sum(dim=1) - 1
-                pooled = hidden[torch.arange(hidden.shape[0]), seq_lens.clamp(min=0)]
-
-                # 3. Нормализация
-                norm = self.text_encoder.text_model.final_layer_norm
-                hidden = norm(hidden)
-                pooled = norm(pooled)
-
-                # 4. Объединяем в матрицу: Пулед (как 1-й токен) + остальные токены
-                # pooled.unsqueeze(1) делает [Batch, 1, Dim]
-                embeds = torch.cat([pooled.unsqueeze(1), hidden], dim=1)
-
-                # 5. Расширяем маску для нового токена (добавляем единицы спереди)
-                ones = torch.ones((toks.attention_mask.shape[0], 1), dtype=toks.attention_mask.dtype, device=device)
-                mask = torch.cat([ones, toks.attention_mask], dim=1)
-
-                return embeds, mask
-                
-        def get_pooled_encode(texts):
-            if texts is None:
-                texts = ""
-
-            if isinstance(texts, str):
-                texts = [texts]
-
-            with torch.no_grad():
-                # 1. Собираем текстовые промпты оборачивая их в Chat Template
-                formatted_prompts = []
-                for t in texts:
-                    messages = [{"role": "user", "content": [{"type": "text", "text": t}]}]
-                    res_text = self.tokenizer2.apply_chat_template(
-                        messages, 
-                        add_generation_prompt=True, 
-                        tokenize=False 
-                    )
-                    formatted_prompts.append(res_text)
-
-                # 2. Токенизируем, режем и добавляем паддинг за один раз
-                toks = self.tokenizer2(
-                    formatted_prompts,
-                    padding="max_length",
-                    max_length=self.text_encoder.config.max_position_embeddings,
-                    truncation=True, # Не забываем обрезать, если вдруг длиннее
-                    return_tensors="pt"
-                ).to(device)
-
-                # 3. Прогоняем через модель
-                outputs = self.text_encoder2(
-                    input_ids=toks.input_ids, 
-                    attention_mask=toks.attention_mask,
-                    output_hidden_states=True
-                )
-
-                layer_index = -2
-                last_hidden = outputs.hidden_states[layer_index]
-                seq_len = toks.attention_mask.sum(dim=1) - 1
-                pooled = last_hidden[torch.arange(len(last_hidden)), seq_len.clamp(min=0)]
-
-                return pooled
-
-        pos_embeds, pos_mask = get_single_encode(prompt)
-        neg_embeds, neg_mask = get_single_encode(negative_prompt)
-        pos_pooled = get_pooled_encode(prompt)
-        neg_pooled = get_pooled_encode(negative_prompt)
-
-        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)
-            neg_pooled = neg_pooled.repeat(batch_size, 1)
-
-        if pos_pooled.shape[0] != batch_size:
-            pos_pooled = pos_pooled.repeat(batch_size, 1)
-
-        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
-        final_mask = torch.cat([neg_mask, pos_mask], dim=0)
-        pooled_embeds = torch.cat([neg_pooled, pos_pooled], dim=0)
-
-        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64), pooled_embeds.to(dtype=dtype)
-
-    @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,
-        refine_prompt: bool = False, # Флаг рефайна!
-        # 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)
-            
-        # ==================== REFINE PROMPT (INLINE) ====================
-        if refine_prompt and prompt:
-            sys_msg = (
-                "You are a skilled text-to-image prompt engineer whose sole function is to transform the user's input into an aesthetically optimized, detailed, and visually descriptive three-sentence output. "
-                "**The primary subject (e.g., 'girl', 'dog', 'house') MUST be the main focus of the revised prompt and MUST be described in rich detail within the first sentence or two.** "
-                "Output **only** the final revised prompt in **English**, with absolutely no commentary.\n Don't use cliches like warm,soft,vibrant, wildflowers. Be creative "
-                "User input prompt: "
-            )
-            prompts_list = [prompt] if isinstance(prompt, str) else prompt
-            refined_list = []
-            
-            for p in prompts_list:
-                messages = [{"role": "user", "content": [{"type": "text", "text": sys_msg + p}]}]
-                
-                # Используем Qwen-Instruct формат (apply_chat_template сам подставит system/user/assistant токены)
-                inputs = self.tokenizer2.apply_chat_template(
-                    messages,
-                    tokenize=True,
-                    add_generation_prompt=True,
-                    return_dict=True,
-                    return_tensors="pt"
-                ).to(device)
-                
-                generated_ids = self.text_encoder2.generate(
-                    **inputs, max_new_tokens=self.text_encoder.config.max_position_embeddings, do_sample=True,temperature = 0.7
-                )
-                
-                # Обрезаем входные токены из ответа
-                generated_ids_trimmed = [
-                    out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-                ]
-                output_text = self.tokenizer2.batch_decode(
-                    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-                )
-                refined_list.append(output_text)
-                
-            prompt = refined_list[0] if isinstance(prompt, str) else refined_list
-
-        # ==================== ENCODE PROMPTS ====================
-        text_embeddings, attention_mask, pooled_embeds = 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:]
-    
-        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
-
-            added_cond_kwargs = {
-                "text_embeds": pooled_embeds,
-            }
-    
-            model_out = self.unet(
-                latent_model_input,
-                t,
-                encoder_hidden_states=text_embeddings,
-                encoder_attention_mask=attention_mask,
-                added_cond_kwargs=added_cond_kwargs,
-                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":
-            if not return_dict: return (latents, prompt)
-            return SdxsPipelineOutput(images=latents, prompt=prompt)
-    
-        latents = latents * vae_scaling_factor + vae_shift_factor
-        latents = self.flux_decode(latents)
-
-        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, prompt)
-        return SdxsPipelineOutput(images=images, prompt=prompt)

pipeline_sdxs-Copy2.py

@@ -1,377 +0,0 @@
-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
-from transformers import Qwen3ForCausalLM, Qwen2Tokenizer
-
-@dataclass
-class SdxsPipelineOutput(BaseOutput):
-    images: Union[List[Image.Image], np.ndarray]
-    prompt: Optional[Union[str, List[str]]] = None  # Возврат улучшенного промпта
-
-class SdxsPipeline(DiffusionPipeline):
-    def __init__(self, vae, text_encoder, text_encoder2, tokenizer, tokenizer2, unet, scheduler):
-        super().__init__()
-        self.register_modules(
-            vae=vae, 
-            text_encoder=text_encoder,
-            text_encoder2=text_encoder2, 
-            tokenizer=tokenizer,
-            tokenizer2=tokenizer2,
-            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)
-        target_width = ((width // self.vae_scale_factor) * self.vae_scale_factor)
-        
-        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]
-
-    @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
-        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):
-            if not texts:
-                texts = [""]
-            elif 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
-                )
-
-                # 1. Берем -2 слой [Batch, Seq, Dim]
-                hidden = outputs.hidden_states[-2]
-
-                # 2. Достаем pooled вектор (последний токен) [Batch, Dim]
-                seq_lens = toks.attention_mask.sum(dim=1) - 1
-                pooled = hidden[torch.arange(hidden.shape[0]), seq_lens.clamp(min=0)]
-
-                # 3. Нормализация
-                norm = self.text_encoder.text_model.final_layer_norm
-                hidden = norm(hidden)
-                pooled = norm(pooled)
-
-                # 4. Объединяем в матрицу: Пулед (как 1-й токен) + остальные токены
-                # pooled.unsqueeze(1) делает [Batch, 1, Dim]
-                embeds = torch.cat([pooled.unsqueeze(1), hidden], dim=1)
-
-                # 5. Расширяем маску для нового токена (добавляем единицы спереди)
-                ones = torch.ones((toks.attention_mask.shape[0], 1), dtype=toks.attention_mask.dtype, device=device)
-                mask = torch.cat([ones, toks.attention_mask], dim=1)
-
-                return embeds, mask, pooled
-                
-        def get_pooled_encode(texts):
-            if texts is None:
-                texts = ""
-
-            if isinstance(texts, str):
-                texts = [texts]
-
-            with torch.no_grad():
-                # 1. Собираем текстовые промпты оборачивая их в Chat Template
-                formatted_prompts = []
-                for t in texts:
-                    messages = [{"role": "user", "content": [{"type": "text", "text": t}]}]
-                    res_text = self.tokenizer2.apply_chat_template(
-                        messages, 
-                        add_generation_prompt=True, 
-                        tokenize=False 
-                    )
-                    formatted_prompts.append(res_text)
-
-                # 2. Токенизируем, режем и добавляем паддинг за один раз
-                toks = self.tokenizer2(
-                    formatted_prompts,
-                    padding="max_length",
-                    max_length=self.text_encoder.config.max_position_embeddings,
-                    truncation=True, # Не забываем обрезать, если вдруг длиннее
-                    return_tensors="pt"
-                ).to(device)
-
-                # 3. Прогоняем через модель
-                outputs = self.text_encoder2(
-                    input_ids=toks.input_ids, 
-                    attention_mask=toks.attention_mask,
-                    output_hidden_states=True
-                )
-
-                layer_index = -2
-                last_hidden = outputs.hidden_states[layer_index]
-                seq_len = toks.attention_mask.sum(dim=1) - 1
-                pooled = last_hidden[torch.arange(len(last_hidden)), seq_len.clamp(min=0)]
-
-                return pooled
-
-        pos_embeds, pos_mask, pooled_pos = get_single_encode(prompt)
-        neg_embeds, neg_mask, pooled_neg = get_single_encode(negative_prompt)
-        # 768 + 2048
-        pos_pooled = torch.cat([pooled_pos, get_pooled_encode(prompt)], dim=1)
-        neg_pooled = torch.cat([pooled_neg, get_pooled_encode(negative_prompt)], dim=1)
-
-        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)
-            neg_pooled = neg_pooled.repeat(batch_size, 1)
-
-        if pos_pooled.shape[0] != batch_size:
-            pos_pooled = pos_pooled.repeat(batch_size, 1)
-
-        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
-        final_mask = torch.cat([neg_mask, pos_mask], dim=0)
-        pooled_embeds = torch.cat([neg_pooled, pos_pooled], dim=0)
-
-        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64), pooled_embeds.to(dtype=dtype)
-
-    @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,
-        refine_prompt: bool = False, # Флаг рефайна!
-        # 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)
-            
-        # ==================== REFINE PROMPT (INLINE) ====================
-        if refine_prompt and prompt:
-            sys_msg = (
-                "You are a skilled text-to-image prompt engineer whose sole function is to transform the user's input into an aesthetically optimized, detailed, and visually descriptive three-sentence output. "
-                "**The primary subject (e.g., 'girl', 'dog', 'house') MUST be the main focus of the revised prompt and MUST be described in rich detail within the first sentence or two.** "
-                "Output **only** the final revised prompt in **English**, with absolutely no commentary.\n Don't use cliches like warm,soft,vibrant, wildflowers. Be creative "
-                "User input prompt: "
-            )
-            prompts_list = [prompt] if isinstance(prompt, str) else prompt
-            refined_list = []
-            
-            for p in prompts_list:
-                messages = [{"role": "user", "content": [{"type": "text", "text": sys_msg + p}]}]
-                
-                # Используем Qwen-Instruct формат (apply_chat_template сам подставит system/user/assistant токены)
-                inputs = self.tokenizer2.apply_chat_template(
-                    messages,
-                    tokenize=True,
-                    add_generation_prompt=True,
-                    return_dict=True,
-                    return_tensors="pt"
-                ).to(device)
-                
-                generated_ids = self.text_encoder2.generate(
-                    **inputs, max_new_tokens=self.text_encoder.config.max_position_embeddings, do_sample=True,temperature = 0.7
-                )
-                
-                # Обрезаем входные токены из ответа
-                generated_ids_trimmed = [
-                    out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-                ]
-                output_text = self.tokenizer2.batch_decode(
-                    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-                )
-                refined_list.append(output_text)
-                
-            prompt = refined_list[0] if isinstance(prompt, str) else refined_list
-
-        # ==================== ENCODE PROMPTS ====================
-        text_embeddings, attention_mask, pooled_embeds = 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
-
-        # ==================== TIME IDS =======================================
-        # time_ids должен иметь ТОТ ЖЕ batch-размер, что и pooled_embeds и text_embeddings
-        # (в твоём encode_prompt они всегда удваиваются из-за CFG)
-        time_ids = torch.zeros(
-            pooled_embeds.shape[0],   # ← вот это главное
-            6, 
-            device=device, 
-            dtype=torch.long
-        )
-    
-        # ==================== 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:]
-    
-        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,
-                #added_cond_kwargs=added_cond_kwargs,
-                added_cond_kwargs={"text_embeds": pooled_embeds,"time_ids": time_ids},
-                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":
-            if not return_dict: return (latents, prompt)
-            return SdxsPipelineOutput(images=latents, prompt=prompt)
-    
-        latents = latents * vae_scaling_factor + vae_shift_factor
-        latents = self.flux_decode(latents)
-
-        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, prompt)
-        return SdxsPipelineOutput(images=images, prompt=prompt)

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:95269c824348ba2d5212db7d149efc938a519b90476a8f172936927d7867b7d6
-size 5599938
+oid sha256:1da01a71ae7db2ab06541fc09ddfd80a8b81f804ab2194f4defdabda9c060f60
+size 4347682

train-Copy1.py

@@ -1,818 +0,0 @@
-#from comet_ml import Experiment
-import os
-os.environ["NCCL_P2P_DISABLE"] = "1"
-# disable this on old GPU?
-os.environ["NCCL_IB_DISABLE"] = "1"
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-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,AutoencoderKLFlux2,FlowMatchEulerDiscreteScheduler
-from accelerate import Accelerator, DeepSpeedPlugin
-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, Qwen2Tokenizer,Qwen3VLForConditionalGeneration
-
-# --------------------------- Параметры ---------------------------
-ds_path = "/workspace/sdxs-1b/datasets/ds1234_flux32"
-project = "unet"
-## total batch (split // num `GPU)
-batch_size = 48
-base_learning_rate = 6e-6
-min_learning_rate = 7e-7
-num_epochs = 8
-sample_interval_share = 5
-cfg_dropout = 0.10
-max_length = 248
-use_wandb = False
-use_comet_ml = False
-save_model = False
-use_decay = True
-fbp = False
-optimizer_type = "adam8bit"
-torch_compile = False
-unet_gradient = True
-loss_normalize = False
-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
-# Включение Flash Attention 2/SDPA #MAX_JOBS=4 pip install flash-attn --no-build-isolation
-torch.backends.cuda.enable_flash_sdp(True)
-torch.backends.cuda.enable_mem_efficient_sdp(True)
-torch.backends.cuda.enable_math_sdp(False) # Отключаем медленный вариант
-save_barrier = 1.25
-warmup_percent = 0.03
-#percentile_clipping = 95
-betta2 = 0.995
-eps = 1e-7
-clip_grad_norm = 1.0
-limit = 0
-checkpoints_folder = ""
-gradient_accumulation_steps = 1
-dtype = torch.float32
-mixed_precision = "no"
-
-# Параметры для диффузии
-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")) + 10000001
-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)
-
-accelerator = Accelerator(
-    mixed_precision=mixed_precision,
-    gradient_accumulation_steps=gradient_accumulation_steps
-)
-device = accelerator.device
-
-print("init")
-
-# --------------------------- Инициализация WandB ---------------------------
-if accelerator.is_main_process:
-    if use_wandb:
-        wandb.init(project=project, 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)
-
-# --------------------------- Загрузка моделей ---------------------------
-#vae = AutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to("cpu").eval()
-#vae = AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-dev",subfolder="vae",torch_dtype=dtype).to(device).eval()
-#vae = AsymmetricAutoencoderKL.from_pretrained("vae",torch_dtype=dtype).to(device).eval()
-vae = AutoencoderKLFlux2.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
-tokenizer = AutoTokenizer.from_pretrained("tokenizer")
-text_encoder = AutoModel.from_pretrained("text_encoder", torch_dtype=torch.float16).to(device).eval()
-tokenizer2 = Qwen2Tokenizer.from_pretrained("tokenizer2")
-text_encoder2 = Qwen3VLForConditionalGeneration.from_pretrained("text_encoder2", torch_dtype=torch.float16).to(device).eval()
-scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("scheduler")
-
-def encode_texts(texts, max_length=max_length):
-    if texts is None:
-        texts = [""]
-    if isinstance(texts, str):
-        texts = [texts]
-
-    with torch.no_grad():
-        # --- 1. CLIP Энкодер ---
-        toks = tokenizer(
-            texts,
-            padding="max_length",
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt"
-        ).to(device)
-
-        outputs = text_encoder(
-            input_ids=toks.input_ids, 
-            attention_mask=toks.attention_mask,
-            output_hidden_states=True
-        )
-
-        # 1. Берем -2 слой [Batch, Seq, Dim]
-        hidden = outputs.hidden_states[-2]
-
-        # 2. Достаем pooled вектор (последний токен) [Batch, Dim]
-        seq_lens = toks.attention_mask.sum(dim=1) - 1
-        pooled_clip = hidden[torch.arange(hidden.shape[0]), seq_lens.clamp(min=0)]
-
-        # 3. Нормализация
-        norm = text_encoder.text_model.final_layer_norm
-        hidden = norm(hidden)
-        pooled_clip = norm(pooled_clip)
-
-        # 4. Объединяем в матрицу: Пулед (как 1-й токен) + остальные токены
-        # pooled.unsqueeze(1) делает [Batch, 1, Dim]
-        prompt_embeds = torch.cat([pooled_clip.unsqueeze(1), hidden], dim=1)
-
-        # 5. Расширяем маску для нового токена (добавляем единицы спереди)
-        ones = torch.ones((toks.attention_mask.shape[0], 1), dtype=toks.attention_mask.dtype, device=device)
-        mask = torch.cat([ones, toks.attention_mask], dim=1)
-
-        # --- 2. QWEN Энкодер (через Chat Template) ---
-        # 1. Собираем текстовые промпты оборачивая их в Chat Template
-        formatted_prompts = []
-        for t in texts:
-            messages = [{"role": "user", "content": [{"type": "text", "text": t}]}]
-            res_text = tokenizer2.apply_chat_template(
-                messages, 
-                add_generation_prompt=True, 
-                tokenize=False 
-            )
-            formatted_prompts.append(res_text)
-
-        # 2. Токенизируем, режем и добавляем паддинг за один раз
-        toks = tokenizer2(
-            formatted_prompts,
-            padding="max_length",
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt"
-        ).to(device)
-
-        # 3. Прогоняем через модель
-        outputs = text_encoder2(
-            input_ids=toks.input_ids, 
-            attention_mask=toks.attention_mask,
-            output_hidden_states=True
-        )
-
-        layer_index = -2
-        last_hidden = outputs.hidden_states[layer_index]
-        seq_len = toks.attention_mask.sum(dim=1) - 1
-        pooled = last_hidden[torch.arange(len(last_hidden)), seq_len.clamp(min=0)]
-        return prompt_embeds.to(dtype), mask, pooled.to(dtype)
-
-
-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
-
-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
-
-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, pooled = encode_texts(texts)
-        
-        fixed_samples[size] = (latents, embeddings, masks, texts, pooled)
-    
-    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/dataset/animesfw") or path.startswith("/workspace/dataset/d4/animesfw"))  for path in x["image_path"]],
-    batched=True,
-    batch_size=10000, # обрабатываем по 10к строк за раз
-    num_proc=8
-)
-print(f"Осталось примеров после фильтрации: {len(dataset)}")
-
-# --- 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, pooled = encode_texts(texts)
-    
-    # attention_mask от токенизатора уже имеет нужный формат, но на всякий случай приведем к long
-    attention_mask = attention_mask.to(dtype=torch.int64)
-
-    return latents, embeddings, attention_mask, pooled
-
-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}")
-        
-
-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
-        )
-    else:
-        raise ValueError(f"Unknown optimizer: {name}")
-
-if fbp:
-    trainable_params = list(unet.parameters())
-    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:
-    # 1. Сначала замораживаем ВСЕ параметры UNet
-    #unet.requires_grad_(False)
-    
-    # 2. Размораживаем только нужные
-    #trainable_params_names = ["conv_in.weight", "conv_in.bias", "conv_out.weight", "conv_out.bias"]
-    #train_params = []
-    
-    #for name, param in unet.named_parameters():
-    #    if any(target in name for target in trainable_params_names):
-    #        param.requires_grad = True
-    #        train_params.append(param)
-    #        print(f"Обучаемый слой: {name}")
-
-    unet.requires_grad_(True)
-    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, uncond_pooled  = 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)
-    uncond_pooled = uncond_pooled.to(device=device).repeat(batch_size, 1)
-
-    return uncond_emb, uncond_mask, uncond_pooled
-    
-# Получаем негативные (пустые) условия для валидации
-uncond_emb, uncond_mask, uncond_pooled = 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_pooled = 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, sample_pooled) 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)
-            sample_pooled = sample_pooled.to(dtype=dtype, 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)
-
-                    neg_pooled_batch = uncond_pooled[0:1].expand(curr_batch_size, -1)
-                    attention_pooled_batch = torch.cat([neg_pooled_batch, sample_pooled], dim=0)
-
-                else:
-                    latent_model_input = latents
-                    text_embeddings_batch = sample_text_embeddings
-                    attention_mask_batch = sample_mask
-                    attention_pooled_batch = sample_pooled
-
-                added_cond_kwargs = {
-                    "text_embeds": attention_pooled_batch,
-                }
-                # Предсказание с передачей всех условий
-                model_out = original_model(
-                    latent_model_input, 
-                    t, 
-                    encoder_hidden_states=text_embeddings_batch, 
-                    encoder_attention_mask=attention_mask_batch,
-                    added_cond_kwargs=added_cond_kwargs,
-                )
-                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
-
-            latents = current_latents.detach() * scaling_factor + shift_factor
-            latents = flux_decode(vae,latents)
-            decoded = vae.decode(latents.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")
-        try:
-            all_generated_images.clear()
-            all_captions.clear()
-            del all_generated_images, all_captions
-            del latents, current_latents, latent_model_input, flow
-            del decoded, decoded_fp32
-            del sample_latents, sample_text_embeddings, sample_mask, sample_pooled # Копии на GPU
-            del model_out
-        except UnboundLocalError:
-            pass
-            
-        # 3. Синхронизируем CUDA перед очисткой
-        torch.cuda.synchronize()
-        # 4. Теперь чистим кэш аллокатора и вызываем GC
-        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,uncond_pooled), 0)
-accelerator.wait_for_everyone()
-
-def save_checkpoint(unet, variant=""):
-    if accelerator.is_main_process:
-        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}"))
-
-        torch.cuda.synchronize()
-        torch.cuda.empty_cache()
-        gc.collect()
-        #unet = unet.to(dtype=dtype) #TODO: wtf???
-
-# --------------------------- Тренировочный цикл ---------------------------
-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, pooled) 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)
-
-            # 3. Время t (сэмплим, как и раньше, но чуть сжимаем края)
-            u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            t = u * (1 - 2 * 1e-5) + 1e-5  # Теперь t строго в (0.00001 ... 0.99999)
-            # интерполяция между x0 и шумом
-            noisy_latents = (1.0 - t.view(-1, 1, 1, 1)) * latents + t.view(-1, 1, 1, 1) * noise
-            # делаем integer timesteps для UNet
-            timesteps = t.to(torch.float32).mul(999.0)
-            timesteps = timesteps.clamp(0, scheduler.config.num_train_timesteps - 1)
-
-            added_cond_kwargs = {
-                "text_embeds": pooled,
-            }
-            # --- Вызов UNet с маской  ---
-            model_pred = unet(
-                noisy_latents, 
-                timesteps, 
-                encoder_hidden_states=embeddings,
-                encoder_attention_mask=attention_mask,
-                added_cond_kwargs=added_cond_kwargs,
-            ).sample
-            
-            target = noise - latents
-
-            mse_loss = F.mse_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
-
-            if (global_step % 100 == 0) or (global_step % sample_interval == 0):
-                accelerator.wait_for_everyone()
-
-            accelerator.backward(mse_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:
-                    grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
-                    grad = grad_val.float().item() if torch.is_tensor(grad_val) else 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["lr"] = current_lr
-                    log_data["grad"] = grad
-                    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 or global_step==50:
-                        # Передаем tuple (emb, mask) для негатива
-                        if save_model:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask,uncond_pooled), global_step)
-                        elif epoch % 10 == 0:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask,uncond_pooled), 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)
-                        unet.train()
-
-    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("Готово!")

train-Copy2.py

@@ -1,843 +0,0 @@
-#from comet_ml import Experiment
-import os
-os.environ["NCCL_P2P_DISABLE"] = "1"
-# disable this on old GPU?
-os.environ["NCCL_IB_DISABLE"] = "1"
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-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,AutoencoderKLFlux2,FlowMatchEulerDiscreteScheduler
-from accelerate import Accelerator, DeepSpeedPlugin
-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, Qwen2Tokenizer,Qwen3VLForConditionalGeneration
-
-# --------------------------- Параметры ---------------------------
-ds_path = "/workspace/sdxs-1b/datasets/ds1234_flux32"
-project = "unet"
-## total batch (split // num `GPU)
-batch_size = 48
-base_learning_rate = 6e-6
-min_learning_rate = 7e-7
-num_epochs = 8
-sample_interval_share = 10
-cfg_dropout = 0.10
-max_length = 248
-use_wandb = False
-use_comet_ml = True
-save_model = True
-use_decay = True
-fbp = False
-optimizer_type = "adam8bit"
-torch_compile = False
-unet_gradient = True
-loss_normalize = False
-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
-# Включение Flash Attention 2/SDPA #MAX_JOBS=4 pip install flash-attn --no-build-isolation
-torch.backends.cuda.enable_flash_sdp(True)
-torch.backends.cuda.enable_mem_efficient_sdp(True)
-torch.backends.cuda.enable_math_sdp(False) # Отключаем медленный вариант
-save_barrier = 1.25
-warmup_percent = 0.03
-#percentile_clipping = 95
-betta2 = 0.995
-eps = 1e-7
-clip_grad_norm = 1.0
-limit = 0
-checkpoints_folder = ""
-gradient_accumulation_steps = 1
-dtype = torch.float32
-mixed_precision = "no"
-
-# Параметры для диффузии
-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")) + 10000001
-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)
-
-accelerator = Accelerator(
-    mixed_precision=mixed_precision,
-    gradient_accumulation_steps=gradient_accumulation_steps
-)
-device = accelerator.device
-
-print("init")
-
-# --------------------------- Инициализация WandB ---------------------------
-if accelerator.is_main_process:
-    if use_wandb:
-        wandb.init(project=project, 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)
-
-# --------------------------- Загрузка моделей ---------------------------
-#vae = AutoencoderKL.from_pretrained("vae", torch_dtype=dtype).to("cpu").eval()
-#vae = AutoencoderKLFlux2.from_pretrained("black-forest-labs/FLUX.2-dev",subfolder="vae",torch_dtype=dtype).to(device).eval()
-#vae = AsymmetricAutoencoderKL.from_pretrained("vae",torch_dtype=dtype).to(device).eval()
-vae = AutoencoderKLFlux2.from_pretrained("vae", torch_dtype=dtype).to(device).eval()
-tokenizer = AutoTokenizer.from_pretrained("tokenizer")
-text_encoder = AutoModel.from_pretrained("text_encoder", torch_dtype=torch.float16).to(device).eval()
-tokenizer2 = Qwen2Tokenizer.from_pretrained("tokenizer2")
-text_encoder2 = Qwen3VLForConditionalGeneration.from_pretrained("text_encoder2", torch_dtype=torch.float16).to(device).eval()
-scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("scheduler")
-
-def encode_texts(texts, max_length=max_length):
-    if texts is None:
-        texts = [""]
-    if isinstance(texts, str):
-        texts = [texts]
-
-    with torch.no_grad():
-        # --- 1. CLIP Энкодер ---
-        toks = tokenizer(
-            texts,
-            padding="max_length",
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt"
-        ).to(device)
-
-        outputs = text_encoder(
-            input_ids=toks.input_ids, 
-            attention_mask=toks.attention_mask,
-            output_hidden_states=True
-        )
-
-        # 1. Берем -2 слой [Batch, Seq, Dim]
-        hidden = outputs.hidden_states[-2]
-
-        # 2. Достаем pooled вектор (последний токен) [Batch, Dim]
-        seq_lens = toks.attention_mask.sum(dim=1) - 1
-        pooled_clip = hidden[torch.arange(hidden.shape[0]), seq_lens.clamp(min=0)]
-
-        # 3. Нормализация
-        norm = text_encoder.text_model.final_layer_norm
-        hidden = norm(hidden)
-        pooled_clip = norm(pooled_clip)
-
-        # 4. Объединяем в матрицу: Пулед (как 1-й токен) + остальные токены
-        # pooled.unsqueeze(1) делает [Batch, 1, Dim]
-        prompt_embeds = torch.cat([pooled_clip.unsqueeze(1), hidden], dim=1)
-
-        # 5. Расширяем маску для нового токена (добавляем единицы спереди)
-        ones = torch.ones((toks.attention_mask.shape[0], 1), dtype=toks.attention_mask.dtype, device=device)
-        mask = torch.cat([ones, toks.attention_mask], dim=1)
-
-        # --- 2. QWEN Энкодер (через Chat Template) ---
-        # 1. Собираем текстовые промпты оборачивая их в Chat Template
-        formatted_prompts = []
-        for t in texts:
-            messages = [{"role": "user", "content": [{"type": "text", "text": t}]}]
-            res_text = tokenizer2.apply_chat_template(
-                messages, 
-                add_generation_prompt=True, 
-                tokenize=False 
-            )
-            formatted_prompts.append(res_text)
-
-        # 2. Токенизируем, режем и добавляем паддинг за один раз
-        toks = tokenizer2(
-            formatted_prompts,
-            padding="max_length",
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt"
-        ).to(device)
-
-        # 3. Прогоняем через модель
-        outputs = text_encoder2(
-            input_ids=toks.input_ids, 
-            attention_mask=toks.attention_mask,
-            output_hidden_states=True
-        )
-
-        layer_index = -2
-        last_hidden = outputs.hidden_states[layer_index]
-        seq_len = toks.attention_mask.sum(dim=1) - 1
-        pooled = last_hidden[torch.arange(len(last_hidden)), seq_len.clamp(min=0)]
-        pooled = torch.cat([pooled_clip, pooled], dim=1)
-        return prompt_embeds.to(dtype), mask, pooled.to(dtype)
-
-
-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
-
-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
-
-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, pooled = encode_texts(texts)
-        
-        fixed_samples[size] = (latents, embeddings, masks, texts, pooled)
-    
-    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/dataset/animesfw") or path.startswith("/workspace/dataset/d4/animesfw"))  for path in x["image_path"]],
-    batched=True,
-    batch_size=10000, # обрабатываем по 10к строк за раз
-    num_proc=8
-)
-print(f"Осталось примеров после фильтрации: {len(dataset)}")
-
-# --- 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, pooled = encode_texts(texts)
-    
-    # attention_mask от токенизатора уже имеет нужный формат, но на всякий случай приведем к long
-    attention_mask = attention_mask.to(dtype=torch.int64)
-
-    return latents, embeddings, attention_mask, pooled
-
-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}")
-        
-
-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
-        )
-    else:
-        raise ValueError(f"Unknown optimizer: {name}")
-
-if fbp:
-    trainable_params = list(unet.parameters())
-    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:
-    # 1. Сначала замораживаем ВСЕ параметры UNet
-    #unet.requires_grad_(False)
-    
-    # 2. Размораживаем только нужные
-    #trainable_params_names = ["conv_in.weight", "conv_in.bias", "conv_out.weight", "conv_out.bias"]
-    #train_params = []
-    
-    #for name, param in unet.named_parameters():
-    #    if any(target in name for target in trainable_params_names):
-    #        param.requires_grad = True
-    #        train_params.append(param)
-    #        print(f"Обучаемый слой: {name}")
-
-    unet.requires_grad_(True)
-    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, uncond_pooled  = 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)
-    uncond_pooled = uncond_pooled.to(device=device).repeat(batch_size, 1)
-
-    return uncond_emb, uncond_mask, uncond_pooled
-    
-# Получаем негативные (пустые) условия для валидации
-uncond_emb, uncond_mask, uncond_pooled = 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_pooled = 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, sample_pooled) 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)
-            sample_pooled = sample_pooled.to(dtype=dtype, 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)
-
-            time_ids = torch.zeros(
-                sample_pooled.shape[0],   # ← вот это главное
-                6, 
-                device=device, 
-                dtype=torch.long
-            )
-
-            for t in scheduler.timesteps:
-                if guidance_scale != 1:
-                    latent_model_input = torch.cat([latents, latents], dim=0)
-                    
-                    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)
-                    
-                    neg_mask_batch = uncond_mask[0:1].expand(curr_batch_size, -1)
-                    attention_mask_batch = torch.cat([neg_mask_batch, sample_mask], dim=0)
-                    
-                    neg_pooled_batch = uncond_pooled[0:1].expand(curr_batch_size, -1)
-                    pooled_batch = torch.cat([neg_pooled_batch, sample_pooled], dim=0)
-                    
-                    # ← КЛЮЧЕВОЕ ИСПРАВЛЕНИЕ — time_ids под текущий удвоенный батч!
-                    time_ids = torch.zeros(
-                        pooled_batch.shape[0],          # 2 * curr_batch_size при CFG
-                        6,
-                        device=device,
-                        dtype=torch.long
-                    )
-                    
-                else:
-                    latent_model_input = latents
-                    text_embeddings_batch = sample_text_embeddings
-                    attention_mask_batch = sample_mask
-                    pooled_batch = sample_pooled
-                    
-                    time_ids = torch.zeros(
-                        pooled_batch.shape[0],
-                        6,
-                        device=device,
-                        dtype=torch.long
-                    )
-            
-                # Теперь всё имеет одинаковый batch size
-                model_out = original_model(
-                    latent_model_input,
-                    t,
-                    encoder_hidden_states=text_embeddings_batch,
-                    encoder_attention_mask=attention_mask_batch,
-                    added_cond_kwargs={
-                        "text_embeds": pooled_batch,
-                        "time_ids": time_ids
-                    },
-                )
-
-                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
-
-            latents = current_latents.detach() * scaling_factor + shift_factor
-            latents = flux_decode(vae,latents)
-            decoded = vae.decode(latents.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")
-        try:
-            all_generated_images.clear()
-            all_captions.clear()
-            del all_generated_images, all_captions
-            del latents, current_latents, latent_model_input, flow
-            del decoded, decoded_fp32
-            del sample_latents, sample_text_embeddings, sample_mask, sample_pooled # Копии на GPU
-            del model_out
-        except UnboundLocalError:
-            pass
-            
-        # 3. Синхронизируем CUDA перед очисткой
-        torch.cuda.synchronize()
-        # 4. Теперь чистим кэш аллокатора и вызываем GC
-        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,uncond_pooled), 0)
-accelerator.wait_for_everyone()
-
-def save_checkpoint(unet, variant=""):
-    if accelerator.is_main_process:
-        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}"))
-
-        torch.cuda.synchronize()
-        torch.cuda.empty_cache()
-        gc.collect()
-        #unet = unet.to(dtype=dtype) #TODO: wtf???
-
-# --------------------------- Тренировочный цикл ---------------------------
-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, pooled) 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)
-
-            # 3. Время t (сэмплим, как и раньше, но чуть сжимаем края)
-            u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            t = u * (1 - 2 * 1e-5) + 1e-5  # Теперь t строго в (0.00001 ... 0.99999)
-            # интерполяция между x0 и шумом
-            noisy_latents = (1.0 - t.view(-1, 1, 1, 1)) * latents + t.view(-1, 1, 1, 1) * noise
-            # делаем integer timesteps для UNet
-            timesteps = t.to(torch.float32).mul(999.0)
-            timesteps = timesteps.clamp(0, scheduler.config.num_train_timesteps - 1)
-
-            time_ids = torch.zeros(
-                pooled.shape[0],   # ← вот это главное
-                6, 
-                device=device, 
-                dtype=torch.long
-            )
-
-            # --- Вызов UNet с маской  ---
-            model_pred = unet(
-                noisy_latents, 
-                timesteps, 
-                encoder_hidden_states=embeddings,
-                encoder_attention_mask=attention_mask,
-                added_cond_kwargs={"text_embeds": pooled,"time_ids": time_ids},
-            ).sample
-            
-            target = noise - latents
-
-            mse_loss = F.mse_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
-
-            if (global_step % 100 == 0) or (global_step % sample_interval == 0):
-                accelerator.wait_for_everyone()
-
-            accelerator.backward(mse_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:
-                    grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
-                    grad = grad_val.float().item() if torch.is_tensor(grad_val) else 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["lr"] = current_lr
-                    log_data["grad"] = grad
-                    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 or global_step==50:
-                        # Передаем tuple (emb, mask) для негатива
-                        if save_model:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask,uncond_pooled), global_step)
-                        elif epoch % 10 == 0:
-                            generate_and_save_samples(fixed_samples, (uncond_emb, uncond_mask,uncond_pooled), 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)
-                        unet.train()
-
-    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("Готово!")

{unet0 → unet}/diffusion_pytorch_model.fp16.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a06d8280ec3d97a35a0cb8bdad5f3d308d899f519ddf931a6af3573171f48a56
-size 2973375056
+oid sha256:863dec221a0c36998de31664f313d0600c27ced6cbd742dae7bb8566008d1c0c
+size 2980309336

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:20df58d34704ec4bf5dd388944f0f799e0f758b6d9e48cc472ee473cdb736b88
-size 5960474736
+oid sha256:863dec221a0c36998de31664f313d0600c27ced6cbd742dae7bb8566008d1c0c
+size 2980309336

unet0/config.json

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

unet0/diffusion_pytorch_model.safetensors

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

unet1.5b-2TE-text-Copy1.ipynb

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

unet1.5b-2TE-text-Copy2.ipynb

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