2511

pipeline_sdxs.py

@@ -23,99 +23,61 @@ class SdxsPipeline(DiffusionPipeline):
         self.max_length = max_length
 
     def encode_prompt(self, prompt=None, negative_prompt=None, device=None, dtype=None):
-        """
-        Кодирование промптов в эмбеддинги.
-        Поведение приближено к ручному инференсу:
-         - 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
-
-        # нормализуем входы в списки
+        dtype = dtype or torch.float16  # Явно указываем float16
+    
+        # Преобразуем в списки
         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:
-                    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:
-                pos_inputs = self.tokenizer(
-                    prompt,
-                    return_tensors="pt",
-                    padding="max_length",    # фиксируем длину
-                    truncation=True,
-                    max_length=self.max_length
-                ).to(device)
-                pos_out = self.text_encoder(
-                    pos_inputs.input_ids,
-                    attention_mask=pos_inputs.attention_mask,
-                    output_hidden_states=True
-                )
-                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,
+    
+        # Если промпты не заданы, используем пустые эмбеддинги
+        if prompt is None and negative_prompt is None:
+            hidden_dim = 1024  # Размерность эмбеддинга Qwen3-0.6B
+            seq_len = 150
+            batch_size = 1
+            return torch.zeros((batch_size, seq_len, hidden_dim), dtype=dtype, device=device)
+    
+        # Токенизация с фиксированным max_length=150 и padding="max_length"
+        def encode_texts(texts, max_length=150):
+            with torch.no_grad():
+                toks = self.tokenizer(
+                    texts,
                     return_tensors="pt",
                     padding="max_length",
                     truncation=True,
-                    max_length=self.max_length
+                    max_length=max_length
                 ).to(device)
-                neg_out = self.text_encoder(
-                    neg_inputs.input_ids,
-                    attention_mask=neg_inputs.attention_mask,
-                    output_hidden_states=True
-                )
-                neg_embeddings = neg_out.hidden_states[-1]  # [B, seq_len, dim]
-            else:
-                neg_embeddings = None
-
-            # Если отсутствует 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]
+                outs = self.text_encoder(**toks, output_hidden_states=True)
+                return outs.hidden_states[-1]
+    
+        # Кодируем позитивные и негативные промпты
+        pos_embeddings = encode_texts(prompt) if prompt is not None else None
+        neg_embeddings = encode_texts(negative_prompt) if negative_prompt is not None else None
+    
+        # Выравниваем размеры batch_size
+        batch_size = max(
+            pos_embeddings.shape[0] if pos_embeddings is not None else 0,
+            neg_embeddings.shape[0] if neg_embeddings is not None else 0
+        )
+    
+        # Повторяем эмбеддинги по batch_size
+        if pos_embeddings is not None and pos_embeddings.shape[0] < batch_size:
+            pos_embeddings = pos_embeddings.repeat(batch_size, 1, 1)
+        if neg_embeddings is not None and neg_embeddings.shape[0] < batch_size:
+            neg_embeddings = neg_embeddings.repeat(batch_size, 1, 1)
+    
+        # Конкатенируем для guidance
+        if pos_embeddings is not None and neg_embeddings is not None:
+            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)
 
-        return text_embeddings  # уже на device и dtype правильные
 
     @torch.no_grad()
     def generate_latents(
@@ -129,34 +91,20 @@ class SdxsPipeline(DiffusionPipeline):
         batch_size: int = 1,
         generator=None,
     ):
-        """Генерация латентов. Поведение guidance согласовано с encode_prompt (uncond перед cond)."""
         device = self.device
-        dtype = next(self.unet.parameters()).dtype
-        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:
-            text_embeddings_for_unet = torch.cat([neg_embeds, pos_embeds], dim=0).to(device=device, dtype=dtype)
+        dtype = torch.float16  # Явно указываем float16
+    
+        # Разделяем эмбеддинги на условные и безусловные
+        if guidance_scale > 1:
+            neg_embeds, pos_embeds = text_embeddings.chunk(2)
+            # Повторяем, если batch_size больше
+            if batch_size > pos_embeds.shape[0]:
+                pos_embeds = pos_embeds.repeat(batch_size, 1, 1)
+                neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
+            text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
         else:
-            # если без guidance, просто используем pos
-            text_embeddings_for_unet = pos_embeds.to(device=device, dtype=dtype)
-
-        # Установка timesteps
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-
+            text_embeddings = text_embeddings.repeat(batch_size, 1, 1)
+    
         # Инициализация латентов
         latent_shape = (
             batch_size,
@@ -165,20 +113,21 @@ class SdxsPipeline(DiffusionPipeline):
             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, encoder_hidden_states=text_embeddings_for_unet).sample
-
-            if do_cfg:
+            latent_input = torch.cat([latents, latents], dim=0) if guidance_scale > 1 else latents
+            noise_pred = self.unet(latent_input, t, text_embeddings).sample
+    
+            if guidance_scale > 1:
                 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