from diffusers import DiffusionPipeline import torch import torch.nn as nn import os from diffusers.utils import BaseOutput from dataclasses import dataclass from typing import List, Union, Optional from PIL import Image import numpy as np import json from safetensors.torch import load_file 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__() # Register components self.register_modules( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=scheduler ) # Get the model path, which is either provided directly or from internal dict model_path = None if hasattr(self, '_internal_dict') and self._internal_dict.get('_name_or_path'): model_path = self._internal_dict.get('_name_or_path') # Get device and dtype from existing components device = "cuda" dtype = torch.float16 # Always load text_projector, regardless of whether one was provided projector_path = None # Try to find projector path if model_path and os.path.exists(f"{model_path}/text_projector"): projector_path = f"{model_path}/text_projector" elif os.path.exists("./text_projector"): projector_path = "./text_projector" if projector_path: # Create and load projector try: with open(f"{projector_path}/config.json", "r") as f: projector_config = json.load(f) # Create Linear layer with bias=False self.text_projector = nn.Linear( in_features=projector_config["in_features"], out_features=projector_config["out_features"], bias=False ) # Load the state dict using safetensors self.text_projector.load_state_dict(load_file(f"{projector_path}/model.safetensors")) self.text_projector.to(device=device, dtype=dtype) print(f"Successfully loaded text_projector from {projector_path}",device, dtype) except Exception as e: print(f"Error loading text_projector: {e}") self.vae_scale_factor = 8 def encode_prompt(self, prompt=None, negative_prompt=None, device=None, dtype=None): """Кодирование текстовых промптов в эмбеддинги. Возвращает: - text_embeddings: Тензор эмбеддингов [batch_size, 1, dim] или [2*batch_size, 1, dim] с guidance """ if prompt is None and negative_prompt is None: raise ValueError("Требуется хотя бы один из параметров: prompt или negative_prompt") # Устанавливаем device и dtype device = device or self.device dtype = dtype or next(self.unet.parameters()).dtype with torch.no_grad(): # Обрабатываем позитивный промпт if prompt is not None: if isinstance(prompt, str): prompt = [prompt] text_inputs = self.tokenizer( prompt, return_tensors="pt", padding="max_length", max_length=512, truncation=True ).to(device) # Получаем эмбеддинги outputs = self.text_encoder(text_inputs.input_ids, text_inputs.attention_mask) last_hidden_state = outputs.last_hidden_state.to(device, dtype=dtype) pos_embeddings = self.text_projector(last_hidden_state[:, 0]) # Добавляем размерность для batch processing if pos_embeddings.ndim == 2: pos_embeddings = pos_embeddings.unsqueeze(1) else: # Создаем пустые эмбеддинги, если нет позитивного промпта # (полезно для некоторых сценариев с unconditional generation) batch_size = len(negative_prompt) if isinstance(negative_prompt, list) else 1 pos_embeddings = torch.zeros( batch_size, 1, self.unet.config.cross_attention_dim, device=device, dtype=dtype ) # Обрабатываем негативный промпт if negative_prompt is not None: if isinstance(negative_prompt, str): negative_prompt = [negative_prompt] # Убеждаемся, что размеры негативного и позитивного промптов совпадают if prompt is not None and len(negative_prompt) != len(prompt): neg_batch_size = len(prompt) if len(negative_prompt) == 1: negative_prompt = negative_prompt * neg_batch_size else: negative_prompt = negative_prompt[:neg_batch_size] neg_inputs = self.tokenizer( negative_prompt, return_tensors="pt", padding="max_length", max_length=512, truncation=True ).to(device) neg_outputs = self.text_encoder(neg_inputs.input_ids, neg_inputs.attention_mask) neg_last_hidden_state = neg_outputs.last_hidden_state.to(device, dtype=dtype) neg_embeddings = self.text_projector(neg_last_hidden_state[:, 0]) if neg_embeddings.ndim == 2: neg_embeddings = neg_embeddings.unsqueeze(1) # Объединяем для classifier-free guidance text_embeddings = torch.cat([neg_embeddings, pos_embeddings], dim=0) else: # Если нет негативного промпта, используем нулевые эмбеддинги batch_size = pos_embeddings.shape[0] neg_embeddings = torch.zeros_like(pos_embeddings) text_embeddings = torch.cat([neg_embeddings, pos_embeddings], dim=0) return text_embeddings.to(device=device, dtype=dtype) @torch.no_grad() def generate_latents( self, text_embeddings, height: int = 576, width: int = 576, num_inference_steps: int = 40, guidance_scale: float = 5.0, latent_channels: int = 16, batch_size: int = 1, generator = None, ): """Генерация латентов с использованием эмбеддингов промптов.""" device = self.device dtype = next(self.unet.parameters()).dtype # Проверка размера эмбеддингов do_classifier_free_guidance = guidance_scale > 0 embedding_dim = text_embeddings.shape[0] // 2 if do_classifier_free_guidance else text_embeddings.shape[0] if batch_size > embedding_dim: # Повторяем эмбеддинги до нужного размера батча if do_classifier_free_guidance: neg_embeds, pos_embeds = text_embeddings.chunk(2) neg_embeds = neg_embeds.repeat(batch_size // embedding_dim, 1, 1) pos_embeds = pos_embeds.repeat(batch_size // embedding_dim, 1, 1) text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0) else: text_embeddings = text_embeddings.repeat(batch_size // embedding_dim, 1, 1) # Установка timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) # Инициализация латентов с заданным seed 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="Генерация"): # Подготовка входных данных if do_classifier_free_guidance: latent_input = torch.cat([latents] * 2) else: latent_input = latents latent_input = self.scheduler.scale_model_input(latent_input, t) # Предсказание шума noise_pred = self.unet(latent_input, t, text_embeddings).sample # Применение guidance if do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * ( noise_pred_text - noise_pred_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] else: 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 = 576, width: int = 576, num_inference_steps: int = 40, guidance_scale: float = 5.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 # Устанавливаем генератор с seed для воспроизводимости generator = None if seed is not None: generator = torch.Generator(device=device).manual_seed(seed) # Получаем эмбеддинги, если они не предоставлены 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=prompt, negative_prompt=negative_prompt, device=device ) else: # Убеждаемся, что эмбеддинги на правильном устройстве 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)