2511

pipeline_sdxs-Copy1.py

@@ -1,210 +0,0 @@
-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)

unet/diffusion_pytorch_model.safetensors

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