2602

pipeline_sdxs-Copy2.py

@@ -0,0 +1,369 @@
+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 Qwen3ForCausalLM, Qwen2Tokenizer
+
+@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, text_encoder2, tokenizer, tokenizer2, unet, scheduler):
+        super().__init__()
+        self.register_modules(
+            vae=vae, 
+            text_encoder=text_encoder,
+            text_encoder2=text_encoder2, 
+            tokenizer=tokenizer,
+            tokenizer2=tokenizer2,
+            unet=unet, 
+            scheduler=scheduler
+        )
+        self.vae_scale_factor = 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)
+        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]
+
+    @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
+        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):
+            if not texts:
+                texts = [""]
+            elif 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
+                )
+
+                # 1. Берем -2 слой [Batch, Seq, Dim]
+                hidden = outputs.hidden_states[-2]
+
+                # 2. Достаем pooled вектор (последний токен) [Batch, Dim]
+                seq_lens = toks.attention_mask.sum(dim=1) - 1
+                pooled = hidden[torch.arange(hidden.shape[0]), seq_lens.clamp(min=0)]
+
+                # 3. Нормализация
+                norm = self.text_encoder.text_model.final_layer_norm
+                hidden = norm(hidden)
+                pooled = norm(pooled)
+
+                # 4. Объединяем в матрицу: Пулед (как 1-й токен) + остальные токены
+                # pooled.unsqueeze(1) делает [Batch, 1, Dim]
+                embeds = torch.cat([pooled.unsqueeze(1), hidden], dim=1)
+
+                # 5. Расширяем маску для нового токена (добавляем единицы спереди)
+                ones = torch.ones((toks.attention_mask.shape[0], 1), dtype=toks.attention_mask.dtype, device=device)
+                mask = torch.cat([ones, toks.attention_mask], dim=1)
+
+                return embeds, mask
+                
+        def get_pooled_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.tokenizer2.apply_chat_template(
+                        messages, 
+                        add_generation_prompt=True, 
+                        tokenize=False 
+                    )
+                    formatted_prompts.append(res_text)
+
+                # 2. Токенизируем, режем и добавляем паддинг за один раз
+                toks = self.tokenizer2(
+                    formatted_prompts,
+                    padding="max_length",
+                    max_length=self.text_encoder.config.max_position_embeddings,
+                    truncation=True, # Не забываем обрезать, если вдруг длиннее
+                    return_tensors="pt"
+                ).to(device)
+
+                # 3. Прогоняем через модель
+                outputs = self.text_encoder2(
+                    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 pooled
+
+        pos_embeds, pos_mask = get_single_encode(prompt)
+        neg_embeds, neg_mask = get_single_encode(negative_prompt)
+        pos_pooled = get_pooled_encode(prompt)
+        neg_pooled = get_pooled_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, # Флаг рефайна!
+        # 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)
+            
+        # ==================== 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.tokenizer2.apply_chat_template(
+                    messages,
+                    tokenize=True,
+                    add_generation_prompt=True,
+                    return_dict=True,
+                    return_tensors="pt"
+                ).to(device)
+                
+                generated_ids = self.text_encoder2.generate(
+                    **inputs, max_new_tokens=self.text_encoder.config.max_position_embeddings, 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.tokenizer2.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
+    
+        # ==================== 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:]
+    
+        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
+
+            added_cond_kwargs = {
+                "text_embeds": pooled_embeds,
+            }
+    
+            model_out = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_embeddings,
+                encoder_attention_mask=attention_mask,
+                added_cond_kwargs=added_cond_kwargs,
+                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 * vae_scaling_factor + vae_shift_factor
+        latents = self.flux_decode(latents)
+
+        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)

pipeline_sdxs.py

@@ -263,7 +263,7 @@ class SdxsPipeline(DiffusionPipeline):
                 generated_ids_trimmed = [
                     out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
                 ]
-                output_text = self.tokenizer2.decode(
+                output_text = self.tokenizer2.batch_decode(
                     generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
                 )
                 refined_list.append(output_text)

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8c86a484e526ff5bd01105e5647b1a4d614df601a678acc2012724d6ea7eafe
+oid sha256:62f976ab8efb5653a2427158529b4b013b3b5034da08e15d18deedd0a42c9dd7
 size 5946605448