pipeline_sdxs.py

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 = 16 #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)//2
        target_width = ((width // self.vae_scale_factor) * self.vae_scale_factor)//2
        
        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
        # Достаем параметры из self.vae
        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, 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.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,
        # 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)
    
        # 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. 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:]
    
            #print(f"img2img → strength={coef:.2f}, sigma={sigma:.3f}, steps={len(timesteps)}")
    
        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,
                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":
            return SdxsPipelineOutput(images=latents)
    
        latents = latents * vae_scaling_factor + vae_shift_factor
        latents = self.flux_decode(latents)
        #latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
        #latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
        #    latents.device, latents.dtype
        #)
        #latents = latents * latents_bn_std + latents_bn_mean
        #latents = self._unpatchify_latents(latents)

        image = self.vae.decode(latents, return_dict=False)[0]
        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
        return SdxsPipelineOutput(images=images)