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)