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
from transformers import Qwen3_5ForConditionalGeneration, Qwen3_5Tokenizer

@dataclass
class SdxsPipelineOutput(BaseOutput):
    images: Union[List[Image.Image], np.ndarray]
    prompt: Optional[Union[str, List[str]]] = None

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)            
        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)
        
    def preprocess_image(self, image: Image.Image, width: int, height: int):
        """Ресайз и центрированный кроп изображения для асимметричного VAE."""
        # Для энкодера с масштабом 8
        target_height = ((height // self.vae_scale_factor) * self.vae_scale_factor)
        target_width = ((width // self.vae_scale_factor) * self.vae_scale_factor)
        
        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]


    def encode_prompt(self, prompt, negative_prompt, device, dtype):
        def get_encode(texts):
            if texts is None:
                texts = ""

            if isinstance(texts, str):
                texts = [texts]

            with torch.no_grad():
                # 1. Собираем текстовые промпты оборачивая их в Chat Template
                formatted_prompts = []
                for t in texts:
                    messages = [{"role": "user", "content": [{"type": "text", "text": t}]}]
                    res_text = self.tokenizer.apply_chat_template(
                        messages, 
                        add_generation_prompt=True, 
                        tokenize=False 
                    )
                    formatted_prompts.append(res_text)

                # 2. Токенизируем, режем и добавляем паддинг за один раз
                toks = self.tokenizer(
                    formatted_prompts,
                    padding="max_length",
                    max_length=255,
                    truncation=True, # Не забываем обрезать, если вдруг длиннее
                    return_tensors="pt"
                ).to(device)

                # 3. Прогоняем через модель
                outputs = self.text_encoder(
                    input_ids=toks.input_ids, 
                    attention_mask=toks.attention_mask,
                    output_hidden_states=True
                )

                layer_index = -2
                last_hidden = outputs.hidden_states[layer_index]
                seq_len = toks.attention_mask.sum(dim=1) - 1
                pooled = last_hidden[torch.arange(len(last_hidden)), seq_len.clamp(min=0)]

                return last_hidden, toks.attention_mask, pooled

        pos_embeds, pos_mask, pos_pooled = get_encode(prompt)
        neg_embeds, neg_mask, neg_pooled = get_encode(negative_prompt)

        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)
            neg_pooled = neg_pooled.repeat(batch_size, 1)

        if pos_pooled.shape[0] != batch_size:
            pos_pooled = pos_pooled.repeat(batch_size, 1)

        text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
        final_mask = torch.cat([neg_mask, pos_mask], dim=0)
        pooled_embeds = torch.cat([neg_pooled, pos_pooled], dim=0)

        return text_embeddings.to(dtype=dtype), final_mask.to(dtype=torch.int64), pooled_embeds.to(dtype=dtype)

    @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,
        refine_prompt: bool = False,
        **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)
            
        # ==================== REFINE PROMPT (INLINE) ====================
        if refine_prompt and prompt:
            sys_msg = (
                "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 three-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 or two.** "
                "Output **only** the final revised prompt in **English**, with absolutely no commentary.\n Don't use cliches like warm,soft,vibrant, wildflowers. Be creative "
                "User input prompt: "
            )
            prompts_list = [prompt] if isinstance(prompt, str) else prompt
            refined_list = []
            
            for p in prompts_list:
                messages = [{"role": "user", "content": [{"type": "text", "text": sys_msg + p}]}]
                
                # Используем Qwen-Instruct формат (apply_chat_template сам подставит system/user/assistant токены)
                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=255, do_sample=True,temperature = 0.7
                )
                
                # Обрезаем входные токены из ответа
                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)
                
            prompt = refined_list[0] if isinstance(prompt, str) else refined_list

        # ==================== ENCODE PROMPTS ====================
        text_embeddings, attention_mask, pooled_embeds = 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

        # ==================== TIME IDS =======================================
        time_ids = torch.zeros(
            pooled_embeds.shape[0],
            6, 
            device=device, 
            dtype=torch.long
        )
    
        # ==================== 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)
            latents_clean = (latents_clean - self.vae_latents_mean.to(device, self.vae.dtype)) / self.vae_latents_std.to(device, self.vae.dtype)
            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:]
    
        else:
            # txt2img
            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,
                added_cond_kwargs={"text_embeds": pooled_embeds,"time_ids": time_ids},
                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":
            if not return_dict: return (latents, prompt)
            return SdxsPipelineOutput(images=latents, prompt=prompt)
    
        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, prompt)
        return SdxsPipelineOutput(images=images, prompt=prompt)