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)