2601

pipeline_sdxs.py

@@ -22,15 +22,108 @@ class SdxsPipeline(DiffusionPipeline):
             unet=unet, 
             scheduler=scheduler
         )
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        # Для асимметричного VAE: 8 энкодер, 16 декодер
+        self.encoder_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)  # 8
+        self.decoder_scale_factor = self.encoder_scale_factor * 2  # 16 (если декодер в 2 раза больше)
+        
+    def create_frequency_soft_cutoff_mask(self, height: int, width: int, cutoff_radius: float, 
+                                        transition_width: float = 5.0, device: torch.device = None) -> torch.Tensor:
+        """Создает плавную маску частотного среза для сохранения структуры."""
+        if device is None:
+            device = torch.device('cpu')
+        
+        u = torch.arange(height, device=device)
+        v = torch.arange(width, device=device)
+        u, v = torch.meshgrid(u, v, indexing='ij')
+        
+        center_u, center_v = height // 2, width // 2
+        frequency_radius = torch.sqrt((u - center_u)**2 + (v - center_v)**2)
+        
+        mask = torch.exp(-(frequency_radius - cutoff_radius)**2 / (2 * transition_width**2))
+        mask = torch.where(frequency_radius <= cutoff_radius, torch.ones_like(mask), mask)
+        
+        return mask
+
+    def generate_structured_noise(
+        self,
+        image_latents: torch.Tensor,
+        cutoff_radius: Optional[float] = None,
+        transition_width: float = 2.0,
+        noise_std: float = 1.0,
+    ) -> torch.Tensor:
+        """
+        Генерирует структурированный шум для латентов с сохранением низкочастотной структуры.
+        
+        Args:
+            image_latents: Чистые латенты изображения [B, C, H, W]
+            cutoff_radius: Радиус среза частот (None = авто-расчет на основе coef)
+            transition_width: Ширина плавного перехода
+            noise_std: Стандартное отклонение шума
+            
+        Returns:
+            Структурированный шум с той же размерностью
+        """
+        batch_size, channels, height, width = image_latents.shape
+        device = image_latents.device
+        dtype = image_latents.dtype
+        
+        # Автоматический расчет cutoff_radius если не задан
+        if cutoff_radius is None:
+            # Сохраняем больше низких частот для лучшей структуры
+            max_radius = min(height, width) / 2
+            cutoff_radius = max_radius * 0.7  # Сохраняем 70% низких частот
+        
+        # Создаем частотную маску
+        freq_mask = self.create_frequency_soft_cutoff_mask(
+            height, width, cutoff_radius, transition_width, device
+        )
+        freq_mask = freq_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
+        
+        # Преобразуем латенты в частотную область
+        fft_image = torch.fft.fft2(image_latents, dim=(-2, -1))
+        fft_shifted = torch.fft.fftshift(fft_image, dim=(-2, -1))
+        
+        # Извлекаем фазу изображения
+        image_phase = torch.angle(fft_shifted)
+        
+        # Генерируем гауссовский шум
+        noise = torch.randn_like(image_latents) * noise_std
+        
+        # Преобразуем шум в частотную область
+        fft_noise = torch.fft.fft2(noise, dim=(-2, -1))
+        fft_noise_shifted = torch.fft.fftshift(fft_noise, dim=(-2, -1))
+        
+        # Извлекаем амплитуду шума
+        noise_magnitude = torch.abs(fft_noise_shifted)
+        noise_phase = torch.angle(fft_noise_shifted)
+        
+        # Смешиваем фазы: низкие частоты - фаза изображения, высокие - фаза шума
+        mixed_phase = freq_mask * image_phase + (1 - freq_mask) * noise_phase
+        
+        # Собираем обратно: амплитуда шума + смешанная фаза
+        fft_combined = noise_magnitude * torch.exp(1j * mixed_phase)
+        fft_unshifted = torch.fft.ifftshift(fft_combined, dim=(-2, -1))
+        
+        # Обратное преобразование
+        structured_noise = torch.fft.ifft2(fft_unshifted, dim=(-2, -1))
+        structured_noise = torch.real(structured_noise)
+        
+        # Нормализуем для сохранения статистики гауссовского шума
+        current_std = torch.std(structured_noise)
+        if current_std > 0:
+            structured_noise = structured_noise / current_std * noise_std
+        
+        return structured_noise.to(dtype)
 
-    # --- ВСПОМОГАТЕЛЬНАЯ ФУНКЦИЯ ДЛЯ ПОДГОТОВКИ ИЗОБРАЖЕНИЯ (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
+        """Ресайз и центрированный кроп изображения для асимметричного VAE."""
+        # Для энкодера с масштабом 8
+        target_height = (height // self.encoder_scale_factor) * self.encoder_scale_factor
+        target_width = (width // self.encoder_scale_factor) * self.encoder_scale_factor
+        
         w, h = image.size
-        aspect_ratio = width / height
+        aspect_ratio = target_width / target_height
+        
         if w / h > aspect_ratio:
             new_w = int(h * aspect_ratio)
             left = (w - new_w) // 2
@@ -40,11 +133,11 @@ class SdxsPipeline(DiffusionPipeline):
             top = (h - new_h) // 2
             image = image.crop((0, top, w, top + new_h))
         
-        image = image.resize((width, height ), resample=Image.LANCZOS)
+        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 = image[None].transpose(0, 3, 1, 2)  # [1, C, H, W]
         image = torch.from_numpy(image)
-        return 2.0 * image - 1.0 # В диапазон [-1, 1]
+        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):
@@ -107,6 +200,7 @@ class SdxsPipeline(DiffusionPipeline):
         generator: Optional[torch.Generator] = None,
         output_type: str = "pil",
         return_dict: bool = True,
+        structure_preservation: float = 0.7,  # Новый параметр: сохранение структуры 0-1
         **kwargs,
     ):
         device = self.device
@@ -122,31 +216,18 @@ class SdxsPipeline(DiffusionPipeline):
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps
         
-        # ИСПРАВЛЕННАЯ ЛОГИКА IMG2IMG
+        # 3. Обработка img2img с структурированным шумом
         if image is not None:
-            # Подготовка изображения
+            # Подготовка изображения (только одно изображение, не дублируем)
             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"
-                    )
+                image_tensor = self.preprocess_image(image, width, height).to(
+                    device=device, dtype=self.vae.dtype
+                )
+            else:
+                # Если список, берем первое
+                image_tensor = self.preprocess_image(image[0], width, height).to(
+                    device=device, dtype=self.vae.dtype
+                )
             
             # Кодируем в латенты
             latents_clean = self.vae.encode(image_tensor).latent_dist.sample(generator=generator)
@@ -155,56 +236,66 @@ class SdxsPipeline(DiffusionPipeline):
             latents_clean = (latents_clean - vae_shift_factor) / vae_scaling_factor
             latents_clean = latents_clean.to(dtype=dtype)
             
-            # Определяем начальный таймстеп
-            # 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 в зависимости от направления
+            # Определяем начальный таймстеп для flow matching
+            # coef контролирует силу: 1.0 - максимальное изменение, 0.0 - минимальное
+            t_start = coef
             
             # Находим ближайший шаг в расписании
-            # Некоторые scheduler имеют reversed порядок (от 1 к 0)
-            if timesteps[0] > timesteps[-1]:
-                # Обычный порядок: от шума к данным
+            # Предполагаем, что timesteps идут от 1 (шум) к 0 (данные)
+            if timesteps[0] > timesteps[-1]:  # Обычный порядок: от шума к данным
                 idx = torch.searchsorted(timesteps, t_start, right=True)
-                idx = min(idx, len(timesteps) - 1)
-            else:
-                # Обратный порядок: от данных к шуму
+                idx = min(idx.item(), 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)
+                idx = len(timesteps) - 1 - min(idx.item(), len(timesteps) - 1)
             
             timesteps = timesteps[idx:]
             
-            # Генерируем шум и применя flow matching формулу
-            noise = torch.randn(
-                latents_clean.shape, 
-                generator=generator, 
-                device=device, 
-                dtype=dtype
+            # Генерируем структурированный шум вместо обычного гауссовского
+            # structure_preservation контролирует сколько низких частот сохранить
+            # 1.0 - максимальное сохранение, 0.0 - минимальное
+            
+            # Автоматический расчет cutoff_radius на основе structure_preservation
+            latent_height, latent_width = latents_clean.shape[-2], latents_clean.shape[-1]
+            max_radius = min(latent_height, latent_width) / 2
+            cutoff_radius = max_radius * structure_preservation
+            
+            # Генерируем структурированный шум
+            structured_noise = self.generate_structured_noise(
+                image_latents=latents_clean,
+                cutoff_radius=cutoff_radius,
+                transition_width=2.0,
+                noise_std=1.0
             )
             
+            # Проверяем, что шум имеет правильную статистику
+            current_std = torch.std(structured_noise)
+            if current_std > 0:
+                structured_noise = structured_noise / current_std  # Нормализуем к std=1
+            
             # Берём актуальное время из 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
+            # Применяем формулу flow matching со структурированным шумом
+            latents = (1.0 - actual_t) * latents_clean + actual_t * structured_noise
             
         else:
-            # TXT2IMG: начинаем с чистого шума
+            # TXT2IMG: начинаем с чистого гауссовского шума
             vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
             vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
             
+            # Для декодера с масштабом 16
+            latent_height = height // self.decoder_scale_factor
+            latent_width = width // self.decoder_scale_factor
+            
             latents = torch.randn(
-                (batch_size, self.unet.config.in_channels, 
-                 height // self.vae_scale_factor, 
-                 width // self.vae_scale_factor),
+                (batch_size, self.unet.config.in_channels, latent_height, latent_width),
                 generator=generator, 
                 device=device, 
                 dtype=dtype
             )
         
-        # 3. Denoising Loop
+        # 4. Denoising Loop
         for i, t in enumerate(tqdm(timesteps, desc="Sampling")):
             # CFG preparation
             latent_model_input = torch.cat([latents] * 2) if guidance_scale > 1 else latents
@@ -226,23 +317,44 @@ class SdxsPipeline(DiffusionPipeline):
             # Euler step для flow matching
             latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]
         
-        # 4. Decode
+        # 5. Decode
         if output_type == "latent":
             return SdxsPipelineOutput(images=latents)
         
-        # Масштабируем обратно
+        # Масштабируем обратно для VAE
         latents = latents * vae_scaling_factor + vae_shift_factor
-        image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
+        image_output = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
+        
+        # Исправляем инвертированное изображение - правильная нормализация
+        # VAE обычно выдает в диапазоне [-1, 1] или [0, 1]
+        # Проверяем диапазон и нормализуем правильно
+        image_min = image_output.min()
+        image_max = image_output.max()
         
-        # Пост-процессинг
-        image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        if image_min < -1.5 and image_max > 1.5:  # Вероятно диапазон [-∞, ∞]
+            image_output = torch.tanh(image_output)  # Приводим к [-1, 1]
+        
+        # Нормализуем к [0, 1] для PIL
+        image_output = (image_output.clamp(-1, 1) + 1) / 2
+        
+        # Для асимметричного VAE: выход декодера больше, возможно нужно ресайзить
+        output_height, output_width = image_output.shape[-2:]
+        if output_height != height or output_width != width:
+            # Простой ресайз до нужного размера
+            image_output = torch.nn.functional.interpolate(
+                image_output, size=(height, width), mode='bicubic', align_corners=False
+            )
+        
+        image_output = image_output.clamp(0, 1)
+        image_np = image_output.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]
+            image_np = (image_np * 255).round().astype("uint8")
+            images = [Image.fromarray(img) for img in image_np]
+        else:
+            images = image_np
         
         if not return_dict:
-            return image
+            return images
         
-        return SdxsPipelineOutput(images=image)
+        return SdxsPipelineOutput(images=images)