2511

pipeline_sdxs-Copy1.py

@@ -0,0 +1,210 @@
+from diffusers import DiffusionPipeline
+import torch
+from diffusers.utils import BaseOutput
+from dataclasses import dataclass
+from typing import List, Union, Optional
+from PIL import Image
+import numpy as np
+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, text_projector=None):
+        super().__init__()
+        self.register_modules(
+            vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, 
+            unet=unet, scheduler=scheduler
+        )
+        self.vae_scale_factor = 8
+
+    def encode_prompt(self, prompt=None, negative_prompt=None, device=None, dtype=None):
+        """Кодирование текстовых промптов в эмбеддинги с выравниванием seq_len."""
+        if prompt is None and negative_prompt is None:
+            raise ValueError("Требуется хотя бы один из параметров: prompt или negative_prompt")
+    
+        device = device or self.device
+        dtype = dtype or next(self.unet.parameters()).dtype
+    
+        # Преобразуем в списки
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if isinstance(negative_prompt, str):
+            negative_prompt = [negative_prompt]
+    
+        # Выравнивание размеров позитивных/негативных списков
+        if prompt is not None and negative_prompt is not None:
+            if len(prompt) != len(negative_prompt):
+                if len(negative_prompt) == 1:
+                    negative_prompt = negative_prompt * len(prompt)
+                elif len(prompt) == 1:
+                    prompt = prompt * len(negative_prompt)
+                else:
+                    n = min(len(prompt), len(negative_prompt))
+                    prompt = prompt[:n]
+                    negative_prompt = negative_prompt[:n]
+    
+        with torch.no_grad():
+            # --- Позитивные эмбеддинги ---
+            if prompt is not None:
+                text_inputs = self.tokenizer(
+                    prompt,
+                    return_tensors="pt",
+                    padding=True,       # динамический паддинг
+                    truncation=True,
+                    max_length=512
+                ).to(device)
+                pos_embeddings = self.text_encoder(
+                    text_inputs.input_ids,
+                    attention_mask=text_inputs.attention_mask,
+                    output_hidden_states=True
+                ).hidden_states[-1]  # [batch, seq_len, dim]
+            else:
+                pos_embeddings = None
+    
+            # --- Негативные эмбеддинги ---
+            if negative_prompt is not None:
+                neg_inputs = self.tokenizer(
+                    negative_prompt,
+                    return_tensors="pt",
+                    padding=True,
+                    truncation=True,
+                    max_length=512
+                ).to(device)
+                neg_embeddings = self.text_encoder(
+                    neg_inputs.input_ids,
+                    attention_mask=neg_inputs.attention_mask,
+                    output_hidden_states=True
+                ).hidden_states[-1]  # [batch, seq_len, dim]
+            else:
+                neg_embeddings = None
+    
+            # --- Выравниваем seq_len ---
+            if pos_embeddings is not None and neg_embeddings is not None:
+                max_len = max(pos_embeddings.shape[1], neg_embeddings.shape[1])
+                if pos_embeddings.shape[1] < max_len:
+                    pad = torch.zeros(pos_embeddings.shape[0], max_len - pos_embeddings.shape[1], pos_embeddings.shape[2], device=pos_embeddings.device, dtype=pos_embeddings.dtype)
+                    pos_embeddings = torch.cat([pos_embeddings, pad], dim=1)
+                if neg_embeddings.shape[1] < max_len:
+                    pad = torch.zeros(neg_embeddings.shape[0], max_len - neg_embeddings.shape[1], neg_embeddings.shape[2], device=neg_embeddings.device, dtype=neg_embeddings.dtype)
+                    neg_embeddings = torch.cat([neg_embeddings, pad], dim=1)
+                text_embeddings = torch.cat([neg_embeddings, pos_embeddings], dim=0)
+            elif pos_embeddings is not None:
+                text_embeddings = pos_embeddings
+            else:
+                text_embeddings = neg_embeddings
+    
+        return text_embeddings.to(device=device, dtype=dtype)
+
+
+    @torch.no_grad()
+    def generate_latents(
+        self,
+        text_embeddings,
+        height: int = 640,
+        width: int = 640,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 5.0,
+        latent_channels: int = 16,
+        batch_size: int = 1,
+        generator=None,
+    ):
+        """Генерация латентов с уч��том любого batch_size и guidance."""
+        device = self.device
+        dtype = next(self.unet.parameters()).dtype
+        do_cfg = guidance_scale > 0
+
+        # Разделяем эмбеддинги на условные и безусловные для guidance
+        if do_cfg:
+            neg_embeds, pos_embeds = text_embeddings.chunk(2)
+            # Повторяем, если batch_size больше эмбеддингов
+            if batch_size > pos_embeds.shape[0]:
+                reps = (batch_size + pos_embeds.shape[0] - 1) // pos_embeds.shape[0]
+                pos_embeds = pos_embeds.repeat(reps, 1, 1)[:batch_size]
+                neg_embeds = neg_embeds.repeat(reps, 1, 1)[:batch_size]
+            text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
+        else:
+            if batch_size > text_embeddings.shape[0]:
+                reps = (batch_size + text_embeddings.shape[0] - 1) // text_embeddings.shape[0]
+                text_embeddings = text_embeddings.repeat(reps, 1, 1)[:batch_size]
+
+        # Установка timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        # Инициализация латентов
+        latent_shape = (
+            batch_size,
+            latent_channels,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor
+        )
+        latents = torch.randn(latent_shape, device=device, dtype=dtype, generator=generator)
+
+        # Процесс диффузии
+        for t in tqdm(self.scheduler.timesteps, desc="Генерация"):
+            latent_input = torch.cat([latents, latents], dim=0) if do_cfg else latents
+            noise_pred = self.unet(latent_input, t, text_embeddings).sample
+
+            if do_cfg:
+                noise_uncond, noise_text = noise_pred.chunk(2)
+                noise_pred = noise_uncond + guidance_scale * (noise_text - noise_uncond)
+
+            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+        return latents
+
+    def decode_latents(self, latents, output_type="pil"):
+        """Декодирование латентов в изображения."""
+        latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        with torch.no_grad():
+            images = self.vae.decode(latents).sample
+        images = (images / 2 + 0.5).clamp(0, 1)
+
+        if output_type == "pil":
+            images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+            images = (images * 255).round().astype("uint8")
+            return [Image.fromarray(image) for image in images]
+        return images.cpu().permute(0, 2, 3, 1).float().numpy()
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Optional[Union[str, List[str]]] = None,
+        height: int = 640,
+        width: int = 512,
+        num_inference_steps: int = 40,
+        guidance_scale: float = 4.0,
+        latent_channels: int = 16,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        batch_size: int = 1,
+        seed: Optional[int] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        text_embeddings: Optional[torch.FloatTensor] = None,
+    ):
+        device = self.device
+        generator = torch.Generator(device=device).manual_seed(seed) if seed is not None else None
+
+        if text_embeddings is None:
+            if prompt is None and negative_prompt is None:
+                raise ValueError("Необходимо указать prompt, negative_prompt или text_embeddings")
+            text_embeddings = self.encode_prompt(prompt, negative_prompt, device=device)
+
+        text_embeddings = text_embeddings.to(device)
+        latents = self.generate_latents(
+            text_embeddings=text_embeddings,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            latent_channels=latent_channels,
+            batch_size=batch_size,
+            generator=generator
+        )
+
+        images = self.decode_latents(latents, output_type=output_type)
+        if not return_dict:
+            return images
+        return SdxsPipelineOutput(images=images)

pipeline_sdxs.py

@@ -12,29 +12,39 @@ class SdxsPipelineOutput(BaseOutput):
     images: Union[List[Image.Image], np.ndarray]
 
 class SdxsPipeline(DiffusionPipeline):
-    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler, text_projector=None):
+    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler, text_projector=None, max_length: int = 150):
         super().__init__()
         self.register_modules(
             vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, 
             unet=unet, scheduler=scheduler
         )
+        # совпадает с тем, что вы используете при ручном инференсе
         self.vae_scale_factor = 8
+        self.max_length = max_length
 
     def encode_prompt(self, prompt=None, negative_prompt=None, device=None, dtype=None):
-        """Кодирование текстовых промптов в эмбеддинги с выравниванием seq_len."""
+        """
+        Кодирование промптов в эмбеддинги.
+        Поведение приближено к ручному инференсу:
+         - padding="max_length", truncation=True, max_length=self.max_length
+         - если negative_prompt отсутствует, возвращаем нулевой uncond с нужной формой
+         - возврат: tensor [batch_uncond + batch_cond, seq_len, hidden_dim]
+           где сначала идут uncond, потом cond (чтобы совпадать с concat для guidance)
+        """
         if prompt is None and negative_prompt is None:
             raise ValueError("Требуется хотя бы один из параметров: prompt или negative_prompt")
-    
+
         device = device or self.device
+        # приводим к dtype unet (важно для совместимости)
         dtype = dtype or next(self.unet.parameters()).dtype
-    
-        # Преобразуем в списки
+
+        # нормализуем входы в списки
         if isinstance(prompt, str):
             prompt = [prompt]
         if isinstance(negative_prompt, str):
             negative_prompt = [negative_prompt]
-    
-        # Выравнивание размеров позитивных/негативных списков
+
+        # equalize list lengths: если один из них длины 1, расширяем — как в вашем ручном коде
         if prompt is not None and negative_prompt is not None:
             if len(prompt) != len(negative_prompt):
                 if len(negative_prompt) == 1:
@@ -45,59 +55,67 @@ class SdxsPipeline(DiffusionPipeline):
                     n = min(len(prompt), len(negative_prompt))
                     prompt = prompt[:n]
                     negative_prompt = negative_prompt[:n]
-    
+
         with torch.no_grad():
             # --- Позитивные эмбеддинги ---
             if prompt is not None:
-                text_inputs = self.tokenizer(
+                pos_inputs = self.tokenizer(
                     prompt,
                     return_tensors="pt",
-                    padding=True,       # динамический паддинг
+                    padding="max_length",    # фиксируем длину
                     truncation=True,
-                    max_length=512
+                    max_length=self.max_length
                 ).to(device)
-                pos_embeddings = self.text_encoder(
-                    text_inputs.input_ids,
-                    attention_mask=text_inputs.attention_mask,
+                pos_out = self.text_encoder(
+                    pos_inputs.input_ids,
+                    attention_mask=pos_inputs.attention_mask,
                     output_hidden_states=True
-                ).hidden_states[-1]  # [batch, seq_len, dim]
+                )
+                pos_embeddings = pos_out.hidden_states[-1]  # [B, seq_len, dim]
             else:
                 pos_embeddings = None
-    
+
             # --- Негативные эмбеддинги ---
             if negative_prompt is not None:
                 neg_inputs = self.tokenizer(
                     negative_prompt,
                     return_tensors="pt",
-                    padding=True,
+                    padding="max_length",
                     truncation=True,
-                    max_length=512
+                    max_length=self.max_length
                 ).to(device)
-                neg_embeddings = self.text_encoder(
+                neg_out = self.text_encoder(
                     neg_inputs.input_ids,
                     attention_mask=neg_inputs.attention_mask,
                     output_hidden_states=True
-                ).hidden_states[-1]  # [batch, seq_len, dim]
+                )
+                neg_embeddings = neg_out.hidden_states[-1]  # [B, seq_len, dim]
             else:
                 neg_embeddings = None
-    
-            # --- Выравниваем seq_len ---
-            if pos_embeddings is not None and neg_embeddings is not None:
-                max_len = max(pos_embeddings.shape[1], neg_embeddings.shape[1])
-                if pos_embeddings.shape[1] < max_len:
-                    pad = torch.zeros(pos_embeddings.shape[0], max_len - pos_embeddings.shape[1], pos_embeddings.shape[2], device=pos_embeddings.device, dtype=pos_embeddings.dtype)
-                    pos_embeddings = torch.cat([pos_embeddings, pad], dim=1)
-                if neg_embeddings.shape[1] < max_len:
-                    pad = torch.zeros(neg_embeddings.shape[0], max_len - neg_embeddings.shape[1], neg_embeddings.shape[2], device=neg_embeddings.device, dtype=neg_embeddings.dtype)
-                    neg_embeddings = torch.cat([neg_embeddings, pad], dim=1)
-                text_embeddings = torch.cat([neg_embeddings, pos_embeddings], dim=0)
-            elif pos_embeddings is not None:
-                text_embeddings = pos_embeddings
-            else:
-                text_embeddings = neg_embeddings
-    
-        return text_embeddings.to(device=device, dtype=dtype)
 
+            # Если отсутствует neg_embeddings, создаём нулевой uncond эмбеддинг
+            if neg_embeddings is None and pos_embeddings is not None:
+                b = pos_embeddings.shape[0]
+                seq_len = pos_embeddings.shape[1]
+                hid = pos_embeddings.shape[2]
+                neg_embeddings = torch.zeros((b, seq_len, hid), device=pos_embeddings.device, dtype=pos_embeddings.dtype)
+
+            # Если отсутствует pos_embeddings (маловероятно), создаём нулевой cond
+            if pos_embeddings is None and neg_embeddings is not None:
+                b = neg_embeddings.shape[0]
+                seq_len = neg_embeddings.shape[1]
+                hid = neg_embeddings.shape[2]
+                pos_embeddings = torch.zeros((b, seq_len, hid), device=neg_embeddings.device, dtype=neg_embeddings.dtype)
+
+            # Приводим dtype к нужному (например float16), чтобы совпадало с unet
+            pos_embeddings = pos_embeddings.to(dtype=dtype, device=device)
+            neg_embeddings = neg_embeddings.to(dtype=dtype, device=device)
+
+            # Теперь формируем итоговый тензор: сначала uncond, затем cond
+            # -- если батч >1 и один из них длиной 1, расширим до нужного размера в __call__ / generate_latents
+            text_embeddings = torch.cat([neg_embeddings, pos_embeddings], dim=0)  # -> [B_uncond + B_cond, seq_len, hid]
+
+        return text_embeddings  # уже на device и dtype правильные
 
     @torch.no_grad()
     def generate_latents(
@@ -111,24 +129,30 @@ class SdxsPipeline(DiffusionPipeline):
         batch_size: int = 1,
         generator=None,
     ):
-        """Генерация латентов с учетом любого batch_size и guidance."""
+        """Генерация латентов. Поведение guidance согласовано с encode_prompt (uncond перед cond)."""
         device = self.device
         dtype = next(self.unet.parameters()).dtype
-        do_cfg = guidance_scale > 0
-
-        # Разделяем эмбеддинги на условные и безусловные для guidance
+        do_cfg = guidance_scale > 1e-5  # true если используется guidance
+
+        # text_embeddings: [B_uncond + B_cond, seq_len, hid]
+        # ожидаем, что B_uncond == B_cond == base_batch (или оба равны 1)
+        # разделим пополам по батчу: сначала uncond, затем cond
+        half = text_embeddings.shape[0] // 2
+        neg_embeds = text_embeddings[:half]   # uncond
+        pos_embeds = text_embeddings[half:]   # cond
+
+        # повторяем эмбеддинги, если нужно увеличить batch_size
+        if batch_size > pos_embeds.shape[0]:
+            reps = (batch_size + pos_embeds.shape[0] - 1) // pos_embeds.shape[0]
+            pos_embeds = pos_embeds.repeat(reps, 1, 1)[:batch_size]
+            neg_embeds = neg_embeds.repeat(reps, 1, 1)[:batch_size]
+
+        # для guidance мы собираем [neg, pos] по батчам (concatenate)
         if do_cfg:
-            neg_embeds, pos_embeds = text_embeddings.chunk(2)
-            # Повторяем, если batch_size больше эмбеддингов
-            if batch_size > pos_embeds.shape[0]:
-                reps = (batch_size + pos_embeds.shape[0] - 1) // pos_embeds.shape[0]
-                pos_embeds = pos_embeds.repeat(reps, 1, 1)[:batch_size]
-                neg_embeds = neg_embeds.repeat(reps, 1, 1)[:batch_size]
-            text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
+            text_embeddings_for_unet = torch.cat([neg_embeds, pos_embeds], dim=0).to(device=device, dtype=dtype)
         else:
-            if batch_size > text_embeddings.shape[0]:
-                reps = (batch_size + text_embeddings.shape[0] - 1) // text_embeddings.shape[0]
-                text_embeddings = text_embeddings.repeat(reps, 1, 1)[:batch_size]
+            # если без guidance, просто используем pos
+            text_embeddings_for_unet = pos_embeds.to(device=device, dtype=dtype)
 
         # Установка timesteps
         self.scheduler.set_timesteps(num_inference_steps, device=device)
@@ -145,7 +169,7 @@ class SdxsPipeline(DiffusionPipeline):
         # Процесс диффузии
         for t in tqdm(self.scheduler.timesteps, desc="Генерация"):
             latent_input = torch.cat([latents, latents], dim=0) if do_cfg else latents
-            noise_pred = self.unet(latent_input, t, text_embeddings).sample
+            noise_pred = self.unet(latent_input, t, encoder_hidden_states=text_embeddings_for_unet).sample
 
             if do_cfg:
                 noise_uncond, noise_text = noise_pred.chunk(2)
@@ -190,9 +214,9 @@ class SdxsPipeline(DiffusionPipeline):
         if text_embeddings is None:
             if prompt is None and negative_prompt is None:
                 raise ValueError("Необходимо указать prompt, negative_prompt или text_embeddings")
-            text_embeddings = self.encode_prompt(prompt, negative_prompt, device=device)
+            text_embeddings = self.encode_prompt(prompt, negative_prompt, device=device, dtype=next(self.unet.parameters()).dtype)
 
-        text_embeddings = text_embeddings.to(device)
+        # text_embeddings уже имеет структуру [B_uncond + B_cond, seq_len, hid], dtype и device совместимы
         latents = self.generate_latents(
             text_embeddings=text_embeddings,
             height=height,
@@ -207,4 +231,4 @@ class SdxsPipeline(DiffusionPipeline):
         images = self.decode_latents(latents, output_type=output_type)
         if not return_dict:
             return images
-        return SdxsPipelineOutput(images=images)
+        return SdxsPipelineOutput(images=images)

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7754efea243246c21b73b743ea55055cff9ed385f22d119ee489931185366cf1
-size 8316949
+oid sha256:e0f7ceb281d9d78b8ed0085e763df363b106df049ee6830bc40d84e6a1c25b34
+size 8326857

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3debe901242f713f340604da75758d75924154f8a87d65221eee85a2bcef6f8c
+oid sha256:f51c65967bb570338af3731ea474bbf1d182549ccd33c6136b531a5e383c57e7
 size 6184944280