2601

pipeline_sdxs.py

@@ -22,9 +22,7 @@ class SdxsPipeline(DiffusionPipeline):
             unet=unet, 
             scheduler=scheduler
         )
-        # Для асимметричного VAE: 8 энкодер, 16 декодер
-        self.encoder_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)  # 8
-        self.decoder_scale_factor = self.encoder_scale_factor * 2  # 16 (если декодер в 2 раза больше)
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         
     def create_frequency_soft_cutoff_mask(self, height: int, width: int, cutoff_radius: float, 
                                         transition_width: float = 5.0, device: torch.device = None) -> torch.Tensor:
@@ -118,8 +116,8 @@ class SdxsPipeline(DiffusionPipeline):
     def preprocess_image(self, image: Image.Image, width: int, height: int):
         """Ресайз и центрированный кроп изображения для асимметричного VAE."""
         # Для энкодера с масштабом 8
-        target_height = ((height // self.encoder_scale_factor) * self.encoder_scale_factor)//2
-        target_width = ((width // self.encoder_scale_factor) * self.encoder_scale_factor)//2
+        target_height = ((height // self.vae_scale_factor) * self.vae_scale_factor)//2
+        target_width = ((width // self.vae_scale_factor) * self.vae_scale_factor)//2
         
         w, h = image.size
         aspect_ratio = target_width / target_height
@@ -198,6 +196,7 @@ class SdxsPipeline(DiffusionPipeline):
         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,
         structure_preservation: float = 0.2,  # Новый параметр: сохранение структуры 0-1
@@ -206,6 +205,13 @@ class SdxsPipeline(DiffusionPipeline):
         device = self.device
         dtype = self.unet.dtype
         
+        if generator is None and seed is not None:
+            if torch.cuda.is_available():
+                generator = torch.Generator(device=device)
+            else:
+                generator = torch.Generator()
+            generator.manual_seed(seed)
+        
         # 1. Encode Prompt
         text_embeddings, attention_mask = self.encode_prompt(
             prompt, negative_prompt, device, dtype
@@ -216,9 +222,9 @@ class SdxsPipeline(DiffusionPipeline):
         self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps  # Все 40 шагов
         
-        print(f"Используем ВСЕ шаги: {len(timesteps)}")
-        print(f"Диапазон таймстепов: [{timesteps[0].item():.3f}, {timesteps[-1].item():.3f}]")
-        print(f"Коэффициент смешивания (coef): {coef}")
+        #print(f"Используем ВСЕ шаги: {len(timesteps)}")
+        #print(f"Диапазон таймстепов: [{timesteps[0].item():.3f}, {timesteps[-1].item():.3f}]")
+        #print(f"Коэффициент смешивания (coef): {coef}")
         
         # 3. Обработка img2img с структурированным шумом
         if image is not None:
@@ -272,8 +278,8 @@ class SdxsPipeline(DiffusionPipeline):
             vae_scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
             vae_shift_factor = getattr(self.vae.config, "shift_factor", 0.0)
             
-            latent_height = height // self.decoder_scale_factor
-            latent_width = width // self.decoder_scale_factor
+            latent_height = height // self.vae_scale_factor
+            latent_width = width // self.vae_scale_factor
             
             latents = torch.randn(
                 (batch_size, self.unet.config.in_channels, latent_height, latent_width),
@@ -312,14 +318,6 @@ class SdxsPipeline(DiffusionPipeline):
         latents = latents * vae_scaling_factor + vae_shift_factor
         image_output = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
         
-        # Исправляем инвертированное изображение - правильная нормализация
-        # VAE обычно выдает в диапазоне [-1, 1] или [0, 1]
-        # Проверяем диапазон и нормализуем правильно
-        image_min = image_output.min()
-        image_max = image_output.max()
-        
-        if image_min < -1.5 and image_max > 1.5:  # Вероятно диапазон [-∞, ∞]
-            image_output = torch.tanh(image_output)  # Приводим к [-1, 1]
         
         # Нормализуем к [0, 1] для PIL
         image_output = (image_output.clamp(-1, 1) + 1) / 2

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1433c2a4011e661ee6f50e959a122d8b90fdbbb205a28c884560c3b921cda3ba
-size 4805047
+oid sha256:4421152d5072d3cdf871fb0043fdaf246e0f5bebb26f5ccaa5dd356dfb2ffad1
+size 5386516

unet/config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a14cb265318c2fc495111e70c6f22f9364fe38511040f7fe97530379027afc52
-size 1812
+oid sha256:1a045ebd48be7f677556a522ead90df150527238b5be038d1c88cf40407dba91
+size 1815

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4bbddbf77f86e9e9beb7ac2b91a4c5510f5e3c50eaa8c2b48c46f9ec26db30c5
-size 3438444088
+oid sha256:f518da35f8eb55f8dd01788ceb77f7c47f2e024fbe8c73f6621cfe7435236a87
+size 4408122256

unet_1b.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92ad12839b75b32d0fae2870c2dee334a11b79baece052adbd43b4d32666d503
+size 4408462864