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]
    prompt: Optional[Union[str, List[str]]] = None

class SdxsPipeline(DiffusionPipeline):
    MAX_TEXT_TOKENS = 248
    
    def __init__(self, vae, text_encoder, processor, tokenizer, unet, scheduler):
        super().__init__()
        self.register_modules(
            vae=vae, 
            text_encoder=text_encoder,
            processor=processor,
            tokenizer=tokenizer,
            unet=unet, 
            scheduler=scheduler
        )
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)            
        mean = getattr(self.vae.config, "latents_mean", None)
        std = getattr(self.vae.config, "latents_std", None)
        if mean is not None and std is not None:
            self.vae_latents_std = torch.tensor(std, device=self.unet.device, dtype=self.unet.dtype).view(1, len(std), 1, 1)
            self.vae_latents_mean = torch.tensor(mean, device=self.unet.device, dtype=self.unet.dtype).view(1, len(mean), 1, 1)

    @staticmethod
    def _pad_tensor_to_length(tensor: torch.Tensor, target_len: int, dim: int = 1, pad_value: float = 0) -> torch.Tensor:
        current_len = tensor.shape[dim]
        if current_len >= target_len:
            return tensor
        pad_size = target_len - current_len
        if tensor.dim() == 3: 
            padding = (0, 0, 0, pad_size, 0, 0)
        elif tensor.dim() == 2:
            padding = (0, pad_size, 0, 0)
        else:
            raise ValueError(f"Unsupported tensor dimension: {tensor.dim()}")
        return torch.nn.functional.pad(tensor, padding, value=pad_value)

    
    @torch.no_grad()
    def refine_prompts(
        self, 
        prompts: Union[str, List[str]], 
        system_prompt: Optional[str] = None, 
        temperature: float = 0.7
    ) -> List[str]:
        """
        Refines a list of prompts using the Text Encoder (LLM).
        
        Args:
            prompts: Single prompt string or list of prompts.
            system_prompt: Custom instruction for the LLM. If None, uses default aesthetic enhancer.
            temperature: Sampling temperature for generation.
            
        Returns:
            List of refined prompts.
        """
        device = self.device
        
        # Default system prompt if none provided
        if system_prompt is None:
            system_prompt = (
                "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 two-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.** "
                "Output **only** the final revised prompt, with absolutely no commentary. "
                "Don't use cliches like warm, soft, vibrant, wildflowers. Be creative. User input prompt: "
            )
            
        pad_id = getattr(self.text_encoder.config, "pad_token_id", None) or \
                 getattr(self.text_encoder.config, "eos_token_id", None)
            
        prompts_list = [prompts] if isinstance(prompts, str) else prompts
        refined_list = []
        
        for p in prompts_list:
            # Prepend system prompt to user input
            full_text = system_prompt + p
            messages = [{"role": "user", "content": [{"type": "text", "text": full_text}]}]
            
            inputs = self.tokenizer.apply_chat_template(
                messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
            ).to(device)
            
            generated_ids = self.text_encoder.generate(
                **inputs, 
                max_new_tokens=self.MAX_TEXT_TOKENS, 
                do_sample=True, 
                temperature=temperature, 
                pad_token_id=pad_id
            )
            
            generated_ids_trimmed = [
                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
            ]
            output_text = self.tokenizer.batch_decode(
                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
            )
            refined_list.append(output_text[0])
            
        return refined_list

    @torch.no_grad()
    def encode_text(self, text: Union[str, List[str]]) -> Tuple[torch.Tensor, torch.Tensor]:
        device = self.device
        dtype = self.unet.dtype
        if text is None: text = ""
        if isinstance(text, str): text = [text]

        formatted_prompts = []
        for t in text:
            messages = [{"role": "user", "content": [{"type": "text", "text": t}]}]
            formatted_prompts.append(self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False))

        toks = self.tokenizer(formatted_prompts, padding="max_length", max_length=self.MAX_TEXT_TOKENS, 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)
        
        last_hidden = outputs.hidden_states[-2]

        return last_hidden.to(dtype=dtype), toks.attention_mask.to(dtype=torch.int64)
        
        
    
    @torch.no_grad()
    def encode_image(self, image: Union[Image.Image, str, List[Union[Image.Image, str]]]) -> Tuple[torch.Tensor, torch.Tensor]:
        device = self.device
        dtype = self.unet.dtype
        if isinstance(image, (str, Image.Image)): image = [image]
        batch_size = len(image)

        all_messages = [[{"role": "user", "content": [{"type": "image", "image": img}]}] for img in image]
        formatted_prompts = [self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) for msgs in all_messages]

        inputs = self.processor(text=formatted_prompts, images=image, return_tensors="pt", padding=True, truncation=False).to(device)
        outputs = self.text_encoder(**inputs, output_hidden_states=True)

        last_hidden = outputs.hidden_states[-2]
        seq_lens = inputs.attention_mask.sum(dim=1) - 1
        pooled = last_hidden[torch.arange(batch_size), seq_lens.clamp(min=0)]

        final_embeddings = torch.cat([pooled.unsqueeze(1), last_hidden], dim=1) 
        final_mask = torch.cat([torch.ones((batch_size, 1), device=device, dtype=inputs.attention_mask.dtype), inputs.attention_mask], dim=1)

        return final_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64)
    @torch.no_grad()
    
    
    def encode_text_and_image_naive(self, text: Union[str, List[str]], image: Optional[Union[Image.Image, List[Image.Image], str, List[str]]] = None, scale = 0.5) -> Tuple[torch.Tensor, torch.Tensor]:
        # 1. Получаем текстовый эмбеддинг
        text_embeds, text_mask = self.encode_text(text)
        
        if image is not None:
            if isinstance(image, (str, Image.Image)): 
                image = [image]
            
            # Если картинка одна, а текстов много - размножаем картинку
            if len(image) == 1 and text_embeds.shape[0] > 1: 
                image = image * text_embeds.shape[0]
            
            # --- НАЧАЛО ВСТАВЛЕННОГО КОДА (Логика из encode_image) ---
            device = self.device
            dtype = self.unet.dtype
            batch_size = len(image)

            all_messages = [[{"role": "user", "content": [{"type": "image", "image": img}]}] for img in image]
            formatted_prompts = [self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) for msgs in all_messages]

            inputs = self.processor(text=formatted_prompts, images=image, return_tensors="pt", padding=True, truncation=False).to(device)
            outputs = self.text_encoder(**inputs, output_hidden_states=True)

            # Берем нужный хайден (-2 слой)
            img_hidden_states = outputs.hidden_states[-2]
            # Берем маску attention из процессора
            img_mask = inputs.attention_mask
            # --- КОНЕЦ ВСТАВЛЕННОГО КОДА ---

            # Применяем масштабирование
            if scale != 1.0:
                img_hidden_states = img_hidden_states * scale

            # Приводим маску и типы данных к соответствию с текстом
            img_mask = img_mask.to(text_mask.dtype)
            img_hidden_states = img_hidden_states.to(dtype=dtype)

            # Объединяем текст и последовательность токенов картинки
            final_embeds = torch.cat([text_embeds, img_hidden_states], dim=1)
            final_mask = torch.cat([text_mask, img_mask], dim=1)
            
            return final_embeds, final_mask
            
        return text_embeds, text_mask

    @torch.no_grad()
    def image_upscale(
        self, 
        image: Union[str, Image.Image, List[Union[str, Image.Image]]], 
        batch_size: int = 1
    ) -> List[Image.Image]:
        """
        Upscales images using asymmetric VAE (x2).
        Uses smart batching: processes in parallel if sizes match, else falls back to sequential.
        """
        images = [image] if isinstance(image, (str, Image.Image)) else image
        
        # 1. Preprocess: Load, Handle Alpha, Pad to %8, Normalize
        batch_data = []
        for img in images:
            if isinstance(img, str): img = Image.open(img)
            if img.mode == "RGBA":
                img = Image.alpha_composite(Image.new("RGBA", img.size, (255, 255, 255)), img)
            img = img.convert("RGB")
            
            w, h = img.size
            pw, ph = (8 - w % 8) % 8, (8 - h % 8) % 8
            if pw or ph:
                padded = Image.new("RGB", (w + pw, h + ph), (255, 255, 255))
                padded.paste(img)
                img = padded
            
            t = torch.from_numpy(np.array(img).astype(np.float32) / 127.5 - 1.0).permute(2, 0, 1)
            batch_data.append((t.to(self.device, torch.float16), w, h))

        # 2. Determine Execution Strategy
        # If all shapes are identical, use batch_size. Else fallback to 1.
        unique_shapes = {t.shape for t, _, _ in batch_data}
        step = batch_size if len(unique_shapes) == 1 else 1
        
        output_images = []
        
        # 3. Process Batches
        for i in range(0, len(batch_data), step):
            chunk = batch_data[i : i + step]
            
            # Stack tensors [B, C, H, W]
            tensors = torch.stack([c[0] for c in chunk])
            
            # Encode -> Decode (using mean for deterministic upscale)
            latents = self.vae.encode(tensors).latent_dist.mean
            latents = latents * self.vae_latents_std.to(latents) + self.vae_latents_mean.to(latents)
            decoded = self.vae.decode(latents.to(self.vae.dtype))[0]
            
            # 4. Post-process: Denormalize and Crop
            decoded = (decoded.clamp(-1, 1) + 1) / 2
            for j, tensor in enumerate(decoded):
                w, h = chunk[j][1], chunk[j][2] # Original sizes
                
                # Crop to exact 2x
                arr = tensor.cpu().permute(1, 2, 0).float().numpy()
                arr = arr[:h * 2, :w * 2]
                
                output_images.append(Image.fromarray((arr * 255).astype("uint8")))
                
        return output_images
    
    @torch.no_grad()
    def __call__(
        self,
        prompt: Optional[Union[str, List[str]]] = None,
        negative_prompt: Optional[Union[str, List[str]]] = None,
        prompt_embeds: Optional[torch.Tensor] = None,
        negative_prompt_embeds: Optional[torch.Tensor] = None,
        prompt_attention_mask: Optional[torch.Tensor] = None,
        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
        latents: Optional[torch.Tensor] = None,
        height: int = 1408,
        width: int = 1024,
        num_inference_steps: int = 40,
        guidance_scale: float = 5.0,
        generator: Optional[torch.Generator] = None,
        seed: Optional[int] = None,
        output_type: str = "pil",
        return_dict: bool = True,
        **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)
        
        do_classifier_free_guidance = guidance_scale > 1.0
            
        # 1. Encode Positive
        if prompt_embeds is None:
            if prompt is None: raise ValueError("`prompt` or `prompt_embeds` required.")
            prompt_embeds, prompt_attention_mask = self.encode_text(prompt)
        prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
        prompt_attention_mask = prompt_attention_mask.to(device=device, dtype=torch.int64)
        batch_size = prompt_embeds.shape[0]

        # 2. Encode Negative (only if CFG is enabled)
        if do_classifier_free_guidance:
            if negative_prompt_embeds is None:
                neg_text = negative_prompt if negative_prompt is not None else ("" if isinstance(prompt, str) else [""] * len(prompt))
                negative_prompt_embeds, negative_prompt_attention_mask = self.encode_text(neg_text)
            
            negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device, dtype=torch.int64)

            # Batch size matching
            if negative_prompt_embeds.shape[0] != batch_size:
                negative_prompt_embeds = negative_prompt_embeds.repeat(batch_size, 1, 1)
                negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(batch_size, 1)

            # 3. Align Length (Padding) for Concat
            max_len = max(prompt_embeds.shape[1], negative_prompt_embeds.shape[1])
            prompt_embeds = self._pad_tensor_to_length(prompt_embeds, max_len, dim=1, pad_value=0)
            negative_prompt_embeds = self._pad_tensor_to_length(negative_prompt_embeds, max_len, dim=1, pad_value=0)
            prompt_attention_mask = self._pad_tensor_to_length(prompt_attention_mask, max_len, dim=1, pad_value=0)
            negative_prompt_attention_mask = self._pad_tensor_to_length(negative_prompt_attention_mask, max_len, dim=1, pad_value=0)

            # 4. Concatenate for CFG: [Neg, Pos]
            text_embeddings = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
            attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
        else:
            # If no CFG, we just use positive embeddings as is
            text_embeddings = prompt_embeds
            attention_mask = prompt_attention_mask
    
        # 5. Scheduler & Latents
        self.scheduler.set_timesteps(num_inference_steps, device=device)
        timesteps = self.scheduler.timesteps

        latent_h = height // self.vae_scale_factor
        latent_w = width // self.vae_scale_factor

        if latents is None:
            latents = torch.randn((batch_size, self.unet.config.in_channels, latent_h, latent_w), generator=generator, device=device, dtype=dtype)
        else:
            latents = latents.to(device=device, dtype=dtype)
    
        # 6. Denoising Loop
        for i, t in enumerate(tqdm(timesteps, desc="Sampling")):
            # Duplicate latents only if doing CFG
            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
    
            model_out = self.unet(
                latent_model_input, t,
                encoder_hidden_states=text_embeddings,
                encoder_attention_mask=attention_mask,
                return_dict=False,
            )[0]
    
            # Perform CFG guidance
            if do_classifier_free_guidance:
                flow_uncond, flow_cond = model_out.chunk(2)
                model_out = flow_uncond + guidance_scale * (flow_cond - flow_uncond)
    
            latents = self.scheduler.step(model_out, t, latents, return_dict=False)[0]
    
        # 7. Decode
        if output_type == "latent":
            if not return_dict: return (latents, prompt)
            return SdxsPipelineOutput(images=latents)
    
        latents = latents * self.vae_latents_std.to(device, self.vae.dtype) + self.vae_latents_mean.to(device, self.vae.dtype)
        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)