2601

pipeline_sdxs-Copy1.py

@@ -0,0 +1,176 @@
+import torch
+import numpy as np
+from PIL import Image
+from typing import List, Union, Optional, Tuple
+from dataclasses import dataclass
+
+from diffusers import DiffusionPipeline
+from diffusers.utils import BaseOutput
+from tqdm import tqdm
+
+@dataclass
+class SdxsPipelineOutput(BaseOutput):
+    images: Union[List[Image.Image], np.ndarray]
+
+class SdxsPipeline(DiffusionPipeline):
+    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler):
+        super().__init__()
+        self.register_modules(
+            vae=vae, 
+            text_encoder=text_encoder, 
+            tokenizer=tokenizer, 
+            unet=unet, 
+            scheduler=scheduler
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    def encode_prompt(self, prompt, negative_prompt, device, dtype):
+        """
+        Полное соответствие функции encode_texts и get_negative_embedding из трейна.
+        """
+        def get_single_encode(texts, is_negative=False):
+            if texts is None or texts == "":
+                # Логика get_negative_embedding из трейна
+                hidden_dim = self.text_encoder.config.hidden_size
+                
+                shape = (1, self.text_encoder.config.max_position_embeddings, hidden_dim)
+                # В трейне для негатива: zeros для эмбеддингов и ones для маски
+                emb = torch.zeros(shape, dtype=dtype, device=device)
+                mask = torch.ones((1, self.text_encoder.config.max_position_embeddings), dtype=torch.int64, device=device)
+                return emb, mask
+
+            if isinstance(texts, str):
+                texts = [texts]
+
+            with torch.no_grad():
+                toks = self.tokenizer(
+                    texts,
+                    padding="max_length",
+                    max_length=self.text_encoder.config.max_position_embeddings,
+                    truncation=True,
+                    return_tensors="pt"
+                ).to(device)
+
+                outputs = self.text_encoder(
+                    input_ids=toks.input_ids, 
+                    attention_mask=toks.attention_mask,
+                    output_hidden_states=True
+                )
+
+                # 1. Выбираем нужный слой. 
+                # -1 — это последний блок трансформера
+                # -2 — это предпоследний (стандарт для большинства современных моделей)
+                layer_index = -2
+                prompt_embeds = outputs.hidden_states[layer_index]
+
+                # 2. ДОБАВЛЯЕМ ФИНАЛЬНУЮ НОРМАЛИЗАЦИЮ
+                # В CLIP после всех блоков стоит слой LayerNorm. 
+                # Если мы берем скрытые состояния напрямую, мы "проскакиваем" его.
+                # Нужно применить его вручную:
+                final_layer_norm = self.text_encoder.text_model.final_layer_norm
+                prompt_embeds = final_layer_norm(prompt_embeds)
+
+                return prompt_embeds, toks.attention_mask
+
+        # Получаем эмбеддинги
+        pos_embeds, pos_mask = get_single_encode(prompt)
+        neg_embeds, neg_mask = get_single_encode(negative_prompt, is_negative=True)
+
+        # Выравнивание батча
+        batch_size = pos_embeds.shape[0]
+        if neg_embeds.shape[0] != batch_size:
+            neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
+            neg_mask = neg_mask.repeat(batch_size, 1)
+
+        # Конкатенация для CFG: [Negative, Positive]
+        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
+        final_mask = torch.cat([neg_mask, pos_mask], dim=0)
+
+        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 40, # Как в трейне n_diffusion_steps
+        guidance_scale: float = 4.0,   # Как в трейне
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        device = self.device
+        self.vae.to(device)
+        vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+        vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
+        
+        # 1. Encode Prompt
+        dtype = self.text_encoder.dtype
+        text_embeddings, attention_mask = self.encode_prompt(
+            prompt, negative_prompt, device, dtype
+        )
+
+        # 2. Prepare Latents
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        latent_channels = self.unet.config.in_channels 
+        
+        latents = torch.randn(
+            (batch_size, latent_channels, height // self.vae_scale_factor, width // self.vae_scale_factor),
+            generator=generator,
+            device=device,
+            dtype=dtype
+        )
+
+        # 3. Настройка Flow Matching шедулера
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 4. Denoising Loop
+        for t in tqdm(timesteps, desc="Sampling"):
+            # CFG input
+            latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents
+            
+            # Flow Matching обычно не требует scale_model_input, 
+            # но оставим для совместимости с интерфейсом шедулера
+            if hasattr(self.scheduler, "scale_model_input"):
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # Predict
+            model_out = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                encoder_attention_mask=attention_mask,
+                return_dict=False,
+            )[0]
+
+            # CFG Logic
+            if guidance_scale > 1:
+                flow_uncond, flow_cond = model_out.chunk(2)
+                model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
+
+            # Step (Flow Matching Euler)
+            latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]
+
+        # 5. Decode
+        if output_type == "latent":
+            return SdxsPipelineOutput(images=latents)
+
+        latents = latents * vae_scaling_factor + vae_shift_factor
+        image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
+        
+        # Пост-процессинг
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+
+        if output_type == "pil":
+            image = (image * 255).round().astype("uint8")
+            image = [Image.fromarray(img) for img in image]
+
+        if not return_dict:
+            return image
+
+        return SdxsPipelineOutput(images=image)

pipeline_sdxs.py

@@ -24,17 +24,31 @@ class SdxsPipeline(DiffusionPipeline):
         )
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
 
+    # --- ВСПОМОГАТЕЛЬНАЯ ФУНКЦИЯ ДЛЯ ПОДГОТОВКИ ИЗОБРАЖЕНИЯ (Img2Img) ---
+    def preprocess_image(self, image: Image.Image, width: int, height: int):
+        """Ресайз и центрированный кроп изображения под нужный размер"""
+        w, h = image.size
+        aspect_ratio = width / 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((width, height), resample=Image.LANCZOS)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[None].transpose(0, 3, 1, 2) # [1, C, H, W]
+        image = torch.from_numpy(image)
+        return 2.0 * image - 1.0 # В диапазон [-1, 1]
+
     def encode_prompt(self, prompt, negative_prompt, device, dtype):
-        """
-        Полное соответствие функции encode_texts и get_negative_embedding из трейна.
-        """
         def get_single_encode(texts, is_negative=False):
             if texts is None or texts == "":
-                # Логика get_negative_embedding из трейна
                 hidden_dim = self.text_encoder.config.hidden_size
-                
                 shape = (1, self.text_encoder.config.max_position_embeddings, hidden_dim)
-                # В трейне для негатива: zeros для эмбеддингов и ones для маски
                 emb = torch.zeros(shape, dtype=dtype, device=device)
                 mask = torch.ones((1, self.text_encoder.config.max_position_embeddings), dtype=torch.int64, device=device)
                 return emb, mask
@@ -57,32 +71,21 @@ class SdxsPipeline(DiffusionPipeline):
                     output_hidden_states=True
                 )
 
-                # 1. Выбираем нужный слой. 
-                # -1 — это последний блок трансформера
-                # -2 — это предпоследний (стандарт для большинства современных моделей)
                 layer_index = -2
                 prompt_embeds = outputs.hidden_states[layer_index]
-
-                # 2. ДОБАВЛЯЕМ ФИНАЛЬНУЮ НОРМАЛИЗАЦИЮ
-                # В CLIP после всех блоков стоит слой LayerNorm. 
-                # Если мы берем скрытые состояния напрямую, мы "проскакиваем" его.
-                # Нужно применить его вручную:
                 final_layer_norm = self.text_encoder.text_model.final_layer_norm
                 prompt_embeds = final_layer_norm(prompt_embeds)
 
                 return prompt_embeds, toks.attention_mask
 
-        # Получаем эмбеддинги
         pos_embeds, pos_mask = get_single_encode(prompt)
         neg_embeds, neg_mask = get_single_encode(negative_prompt, is_negative=True)
 
-        # Выравнивание батча
         batch_size = pos_embeds.shape[0]
         if neg_embeds.shape[0] != batch_size:
             neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
             neg_mask = neg_mask.repeat(batch_size, 1)
 
-        # Конкатенация для CFG: [Negative, Positive]
         text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
         final_mask = torch.cat([neg_mask, pos_mask], dim=0)
 
@@ -92,11 +95,13 @@ class SdxsPipeline(DiffusionPipeline):
     def __call__(
         self,
         prompt: Union[str, List[str]],
+        image: Optional[Union[Image.Image, List[Image.Image]]] = None, # Добавлен параметр изображения
+        coef: float = 0.5, # Коэффициент влияния (strength): 1.0 - полный шум, 0.0 - оригинал
         negative_prompt: Optional[Union[str, List[str]]] = None,
         height: int = 1024,
         width: int = 1024,
-        num_inference_steps: int = 40, # Как в трейне n_diffusion_steps
-        guidance_scale: float = 4.0,   # Как в трейне
+        num_inference_steps: int = 40,
+        guidance_scale: float = 4.0,
         generator: Optional[torch.Generator] = None,
         output_type: str = "pil",
         return_dict: bool = True,
@@ -113,28 +118,53 @@ class SdxsPipeline(DiffusionPipeline):
             prompt, negative_prompt, device, dtype
         )
 
-        # 2. Prepare Latents
+        # 2. Определение размеров и батча
         batch_size = 1 if isinstance(prompt, str) else len(prompt)
         latent_channels = self.unet.config.in_channels 
-        
-        latents = torch.randn(
-            (batch_size, latent_channels, height // self.vae_scale_factor, width // self.vae_scale_factor),
-            generator=generator,
-            device=device,
-            dtype=dtype
-        )
 
-        # 3. Настройка Flow Matching шедулера
+        # --- ЛОГИКА ПОДГОТОВКИ ЛАТЕНТОВ (TXT2IMG vs IMG2IMG) ---
+        if image is not None:
+            # Если передана одна картинка на батч текстов - размножаем её
+            if isinstance(image, Image.Image):
+                image = [image] * batch_size
+            
+            # Обработка картинок: ресайз/кроп и перевод в тензоры
+            image_tensors = torch.cat([self.preprocess_image(img, width, height) for img in image]).to(device=device, dtype=vae_dtype := self.vae.dtype)
+            
+            # Кодируем в латенты
+            init_latents = self.vae.encode(image_tensors).latent_dist.sample(generator=generator)
+            init_latents = (init_latents - vae_shift_factor) / vae_scaling_factor
+            init_latents = init_latents.to(dtype=dtype)
+
+            # Создаем шум
+            noise = torch.randn(init_latents.shape, generator=generator, device=device, dtype=dtype)
+
+            # Flow Matching шум: x_t = (1 - t)*x0 + t*epsilon
+            # Здесь t = coef. Если coef=1.0, получаем чистый шум.
+            latents = (1.0 - coef) * init_latents + coef * noise
+        else:
+            # Стандартный txt2img (чистый шум)
+            latents = torch.randn(
+                (batch_size, latent_channels, height // self.vae_scale_factor, width // self.vae_scale_factor),
+                generator=generator,
+                device=device,
+                dtype=dtype
+            )
+
+        # 3. Настройка таймстепов
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps
 
+        # --- ОБРЕЗКА ТАЙМСТЕПОВ ДЛЯ IMG2IMG ---
+        if image is not None:
+            # Оставляем только те шаги, которые соответствуют уровню шума coef
+            # В Flow Matching t идет от 1.0 к 0.0. Нам нужны шаги <= coef.
+            timesteps = timesteps[timesteps <= coef]
+
         # 4. Denoising Loop
         for t in tqdm(timesteps, desc="Sampling"):
-            # CFG input
             latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents
             
-            # Flow Matching обычно не требует scale_model_input, 
-            # но оставим для совместимости с интерфейсом шедулера
             if hasattr(self.scheduler, "scale_model_input"):
                 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
 
@@ -152,7 +182,7 @@ class SdxsPipeline(DiffusionPipeline):
                 flow_uncond, flow_cond = model_out.chunk(2)
                 model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
 
-            # Step (Flow Matching Euler)
+            # Step
             latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]
 
         # 5. Decode

train.py

@@ -1,5 +1,8 @@
 #from comet_ml import Experiment
 import os
+os.environ["NCCL_P2P_DISABLE"] = "1"
+os.environ["NCCL_IB_DISABLE"] = "1"
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 import math
 import torch
 import numpy as np
@@ -9,7 +12,7 @@ from torch.utils.data.distributed import DistributedSampler
 from torch.optim.lr_scheduler import LambdaLR
 from collections import defaultdict
 from diffusers import UNet2DConditionModel, AutoencoderKL,AutoencoderKLFlux2,AsymmetricAutoencoderKL,FlowMatchEulerDiscreteScheduler
-from accelerate import Accelerator
+from accelerate import Accelerator, DeepSpeedPlugin
 from datasets import load_from_disk
 from tqdm import tqdm
 from PIL import Image, ImageOps
@@ -29,11 +32,12 @@ from transformers import AutoTokenizer, AutoModel
 # --------------------------- Параметры ---------------------------
 ds_path = "/workspace/sdxs-08b/datasets/d123_640_sd15"
 project = "unet"
-batch_size = 20 ## total batch (split // num `GPU)
+## total batch (split // num `GPU)
+batch_size = 20
 base_learning_rate = 2e-6
-min_learning_rate = 9e-7
-num_epochs = 3
-sample_interval_share = 5
+min_learning_rate = 7e-7
+num_epochs = 2
+sample_interval_share = 10
 cfg_dropout = 0.10
 max_length = 248
 use_wandb = True
@@ -55,7 +59,6 @@ torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.enable_flash_sdp(True)
 torch.backends.cuda.enable_mem_efficient_sdp(True)
 torch.backends.cuda.enable_math_sdp(False) # Отключаем медленный вариант
-dtype = torch.float32
 save_barrier = 1.05
 warmup_percent = 0.03
 percentile_clipping = 95
@@ -64,17 +67,9 @@ eps = 1e-7
 clip_grad_norm = 1.0
 limit = 0
 checkpoints_folder = ""
-mixed_precision = "no"
 gradient_accumulation_steps = 1
-
-# Это критично для новых карт, чтобы избежать фрагментации при больших контекстах
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:128"
-
-accelerator = Accelerator(
-    mixed_precision=mixed_precision,
-    gradient_accumulation_steps=gradient_accumulation_steps
-)
-device = accelerator.device
+dtype = torch.float32
+mixed_precision = "no"
 
 # Параметры для диффузии
 n_diffusion_steps = 40
@@ -95,6 +90,12 @@ if fixed_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 ---------------------------
@@ -128,7 +129,7 @@ tokenizer = AutoTokenizer.from_pretrained("tokenizer")
 text_model = AutoModel.from_pretrained("text_encoder").to(device).eval()
 scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("scheduler")
 
-def encode_texts(texts, max_length=max_length): # Для SD 1.5 лучше жестко 77
+def encode_texts(texts, max_length=max_length):
     if texts is None:
         texts = [""]
     
@@ -645,7 +646,7 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             if not fbp:
                 if accelerator.sync_gradients:
                     grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
-                    grad = float(grad_val)
+                    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)

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9f32520c8e187f4a57efc07d926f8d2c0078a2e2e80adf5eaf811311aedbecd8
-size 1719263584
+oid sha256:866fefbc206ff908dbc2fd5da4bcc21215fed0e3470c16b494eb4fb3d5810804
+size 3438444088

unet1b.ipynb

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