2601

pipeline_sdxs-Copy2.py

@@ -0,0 +1,196 @@
+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)
+
+    # --- ВСПОМОГАТЕЛЬНАЯ ФУНКЦИЯ ДЛЯ ПОДГОТОВКИ ИЗОБРАЖЕНИЯ (Img2Img) ---
+    def preprocess_image(self, image: Image.Image, width: int, height: int):
+        """Ресайз и центрированный кроп изображения под нужный размер"""
+        height = (height // self.vae_scale_factor) * self.vae_scale_factor // 2
+        width= (width // self.vae_scale_factor) * self.vae_scale_factor // 2
+        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):
+        def get_single_encode(texts, is_negative=False):
+            if texts is None or texts == "":
+                hidden_dim = self.text_encoder.config.hidden_size
+                shape = (1, self.text_encoder.config.max_position_embeddings, hidden_dim)
+                emb = torch.zeros(shape, dtype=dtype, device=device)
+                mask = torch.ones((1, self.text_encoder.config.max_position_embeddings), dtype=torch.int64, device=device)
+                return emb, mask
+
+            if isinstance(texts, str):
+                texts = [texts]
+
+            with torch.no_grad():
+                toks = self.tokenizer(
+                    texts,
+                    padding="max_length",
+                    max_length=self.text_encoder.config.max_position_embeddings,
+                    truncation=True,
+                    return_tensors="pt"
+                ).to(device)
+
+                outputs = self.text_encoder(
+                    input_ids=toks.input_ids, 
+                    attention_mask=toks.attention_mask,
+                    output_hidden_states=True
+                )
+
+                layer_index = -2
+                prompt_embeds = outputs.hidden_states[layer_index]
+                final_layer_norm = self.text_encoder.text_model.final_layer_norm
+                prompt_embeds = final_layer_norm(prompt_embeds)
+
+                return prompt_embeds, toks.attention_mask
+
+        pos_embeds, pos_mask = get_single_encode(prompt)
+        neg_embeds, neg_mask = get_single_encode(negative_prompt, is_negative=True)
+
+        batch_size = pos_embeds.shape[0]
+        if neg_embeds.shape[0] != batch_size:
+            neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
+            neg_mask = neg_mask.repeat(batch_size, 1)
+
+        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
+        final_mask = torch.cat([neg_mask, pos_mask], dim=0)
+
+        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        image: Optional[Union[Image.Image, List[Image.Image]]] = None, # Добавлен параметр изображения
+        coef: float = 0.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,
+        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)
+device = self.device
+        dtype = self.unet.dtype
+        
+        # 1. Encode Prompt
+        text_embeddings, attention_mask = self.encode_prompt(prompt, negative_prompt, device, dtype)
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+
+        # 2. Настройка таймстепов
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps # Обычно от 1.0 до 0.001
+
+        # --- ИСПРАВЛЕННАЯ ЛОГИКА IMG2IMG ---
+        if image is not None:
+            # Превращаем картинку в латенты
+            image_tensor = self.preprocess_image(image, width, height).to(device=device, dtype=self.vae.dtype)
+            
+            # ВАЖНО: Кодируем и применяем масштабирование VAE СРАЗУ
+            latents_clean = self.vae.encode(image_tensor).latent_dist.sample(generator=generator)
+            latents_clean = (latents_clean - vae_shift_factor) / vae_scaling_factor
+            latents_clean = latents_clean.to(dtype=dtype)
+
+            # Определяем начальный таймстеп t_start. 
+            # coef=1.0 -> t=1.0 (шум), coef=0.0 -> t=0.0 (картинка)
+            t_start = coef 
+            
+            # Находим индекс шага в расписании шедулера, который ближе всего к t_start
+            # Это исключает "одинаковость" при разной силе
+            idx = (timesteps - t_start).abs().argmin().item()
+            timesteps = timesteps[idx:] 
+            
+            # Добавляем шум строго по формуле Flow Matching: x_t = (1-t)x_0 + t*eps
+            noise = torch.randn(latents_clean.shape, generator=generator, device=device, dtype=dtype)
+            actual_t = timesteps[0] # Берем точное значение из шедулера
+            latents = (1.0 - actual_t) * latents_clean + actual_t * noise
+        else:
+            # TXT2IMG: Начинаем с чистого шума (t=1.0)
+            latents = torch.randn(
+                (batch_size, self.unet.config.in_channels, height // self.vae_scale_factor, width // self.vae_scale_factor),
+                generator=generator, device=device, dtype=dtype
+            )
+
+        # 3. Denoising Loop
+        for t in tqdm(timesteps, desc="Sampling"):
+            latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents
+            
+            # Predict
+            model_out = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                encoder_attention_mask=attention_mask,
+                return_dict=False,
+            )[0]
+
+            # CFG
+            if guidance_scale > 1:
+                flow_uncond, flow_cond = model_out.chunk(2)
+                model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
+
+            # Step (Эйлеров шаг: latents = latents + (t_next - t_curr) * model_out)
+            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

@@ -97,8 +97,8 @@ 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 - оригинал
+        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,
@@ -110,55 +110,106 @@ class SdxsPipeline(DiffusionPipeline):
         **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)
-device = self.device
         dtype = self.unet.dtype
         
         # 1. Encode Prompt
-        text_embeddings, attention_mask = self.encode_prompt(prompt, negative_prompt, device, dtype)
+        text_embeddings, attention_mask = self.encode_prompt(
+            prompt, negative_prompt, device, dtype
+        )
         batch_size = 1 if isinstance(prompt, str) else len(prompt)
-
+        
         # 2. Настройка таймстепов
         self.scheduler.set_timesteps(num_inference_steps, device=device)
-        timesteps = self.scheduler.timesteps # Обычно от 1.0 до 0.001
-
-        # --- ИСПРАВЛЕННАЯ ЛОГИКА IMG2IMG ---
+        timesteps = self.scheduler.timesteps
+        
+        # ИСПРАВЛЕННАЯ ЛОГИКА IMG2IMG
         if image is not None:
-            # Превращаем картинку в латенты
-            image_tensor = self.preprocess_image(image, width, height).to(device=device, dtype=self.vae.dtype)
+            # Подготовка изображения
+            if isinstance(image, Image.Image):
+                image = [image]
+            
+            # Обработка каждого изображения
+            image_tensors = []
+            for img in image:
+                img_tensor = self.preprocess_image(img, width, height)
+                image_tensors.append(img_tensor)
+            
+            image_tensor = torch.cat(image_tensors, dim=0).to(
+                device=device, dtype=self.vae.dtype
+            )
+            
+            # Проверка batch size
+            if image_tensor.shape[0] != batch_size:
+                if image_tensor.shape[0] == 1:
+                    image_tensor = image_tensor.repeat(batch_size, 1, 1, 1)
+                else:
+                    raise ValueError(
+                        f"Batch size mismatch: prompt has {batch_size} elements, "
+                        f"but image has {image_tensor.shape[0]} elements"
+                    )
             
-            # ВАЖНО: Кодируем и применяем масштабирование VAE СРАЗУ
+            # Кодируем в латенты
             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=dtype)
-
-            # Определяем начальный таймстеп t_start. 
-            # coef=1.0 -> t=1.0 (шум), coef=0.0 -> t=0.0 (картинка)
-            t_start = coef 
-            
-            # Находим индекс шага в расписании шедулера, который ближе всего к t_start
-            # Это исключает "одинаковость" при разной силе
-            idx = (timesteps - t_start).abs().argmin().item()
-            timesteps = timesteps[idx:] 
-            
-            # Добавляем шум строго по формуле Flow Matching: x_t = (1-t)x_0 + t*eps
-            noise = torch.randn(latents_clean.shape, generator=generator, device=device, dtype=dtype)
-            actual_t = timesteps[0] # Берем точное значение из шедулера
+            
+            # Определяем начальный таймстеп
+            # Flow matching: x_t = (1-t)x_0 + t*eps
+            # coef=1.0 -> t=1.0 (полный шум)
+            # coef=0.0 -> t=0.0 (оригинал)
+            
+            # Если scheduler использует [0, 1] диапазон
+            t_start = coef  # или 1.0 - coef в зависимости от направления
+            
+            # Находим ближайший шаг в расписании
+            # Некоторые scheduler имеют reversed порядок (от 1 к 0)
+            if timesteps[0] > timesteps[-1]:
+                # Обычный порядок: от шума к данным
+                idx = torch.searchsorted(timesteps, t_start, right=True)
+                idx = min(idx, len(timesteps) - 1)
+            else:
+                # Обратный порядок: от данных к шуму
+                idx = torch.searchsorted(timesteps.flip(0), 1.0 - t_start, right=True)
+                idx = len(timesteps) - 1 - min(idx, len(timesteps) - 1)
+            
+            timesteps = timesteps[idx:]
+            
+            # Генерируем шум и применя flow matching формулу
+            noise = torch.randn(
+                latents_clean.shape, 
+                generator=generator, 
+                device=device, 
+                dtype=dtype
+            )
+            
+            # Берём актуальное время из scheduler
+            actual_t = timesteps[0].item() if isinstance(timesteps[0], torch.Tensor) else timesteps[0]
+            
+            # Применяем формулу flow matching
             latents = (1.0 - actual_t) * latents_clean + actual_t * noise
+            
         else:
-            # TXT2IMG: Начинаем с чистого шума (t=1.0)
+            # TXT2IMG: начинаем с чистого шума
+            vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+            vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
+            
             latents = torch.randn(
-                (batch_size, self.unet.config.in_channels, height // self.vae_scale_factor, width // self.vae_scale_factor),
-                generator=generator, device=device, dtype=dtype
+                (batch_size, self.unet.config.in_channels, 
+                 height // self.vae_scale_factor, 
+                 width // self.vae_scale_factor),
+                generator=generator, 
+                device=device, 
+                dtype=dtype
             )
-
+        
         # 3. Denoising Loop
-        for t in tqdm(timesteps, desc="Sampling"):
+        for i, t in enumerate(tqdm(timesteps, desc="Sampling")):
+            # CFG preparation
             latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents
             
-            # Predict
+            # Predict flow
             model_out = self.unet(
                 latent_model_input,
                 t,
@@ -166,31 +217,32 @@ device = self.device
                 encoder_attention_mask=attention_mask,
                 return_dict=False,
             )[0]
-
+            
             # CFG
             if guidance_scale > 1:
                 flow_uncond, flow_cond = model_out.chunk(2)
                 model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
-
-            # Step (Эйлеров шаг: latents = latents + (t_next - t_curr) * model_out)
+            
+            # Euler step для flow matching
             latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]
-
-        # 5. Decode
+        
+        # 4. 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)