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)