bf16

pipeline_sdxs.py

@@ -5,36 +5,14 @@ from typing import List, Union, Optional, Tuple
 from dataclasses import dataclass
 import logging
 import math
+import requests
+from io import BytesIO
+import warnings
+from tqdm import tqdm
 
 from transformers import logging as transformers_logging
 from diffusers import DiffusionPipeline
 from diffusers.utils import BaseOutput
-from tqdm import tqdm
-import warnings
-import requests
-from io import BytesIO
-
-def pad_to_match(tensor1, tensor2):
-    """Дополняет нулями короткий тензор, чтобы размерности seq_len совпадали."""
-    len1, len2 = tensor1.shape[1], tensor2.shape[1]
-    if len1 == len2:
-        return tensor1, tensor2
-    if len1 < len2:
-        padding = torch.zeros(tensor1.shape[0], len2 - len1, tensor1.shape[2], device=tensor1.device, dtype=tensor1.dtype)
-        return torch.cat([tensor1, padding], dim=1), tensor2
-    else:
-        padding = torch.zeros(tensor2.shape[0], len1 - len2, tensor2.shape[2], device=tensor2.device, dtype=tensor2.dtype)
-        return tensor1, torch.cat([tensor2, padding], dim=1)
-
-from PIL import Image
-import requests
-from io import BytesIO
-from typing import Union, Tuple
-
-from PIL import Image
-import requests
-from io import BytesIO
-from typing import Union, Tuple
 
 
 @dataclass
@@ -43,7 +21,7 @@ class SdxsPipelineOutput(BaseOutput):
     prompt: Optional[Union[str, List[str]]] = None
 
 class SdxsPipeline(DiffusionPipeline):
-    MAX_TEXT_TOKENS = 300 
+    MAX_TEXT_TOKENS = 250 
     
     def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler):
         super().__init__()
@@ -156,7 +134,7 @@ class SdxsPipeline(DiffusionPipeline):
                 do_sample=True, 
                 temperature=temperature, 
                 top_p=0.95,             # Вместо top_k используем top_p
-                repetition_penalty=1.0, # Снижено с 1.15 (presence_penalty в HF заменяется этим)
+                repetition_penalty=1.15, # Снижено с 1.15 (presence_penalty в HF заменяется этим)
                 pad_token_id=self.tokenizer.pad_token_id,
             )
             
@@ -168,7 +146,7 @@ class SdxsPipeline(DiffusionPipeline):
                 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
             )[0]
             
-            # 3. ВАЖНО: Вырезаем блок размышлений <think&gt;...&lt;/think&gt; из финального ответа
+            # 3. ВАЖНО: Вырезаем блок размышлений <think>...</think> из финального ответа
             # Если модель думала, она выдаст мысль, а потом ответ. Нам нужен только ответ.
             if "</think" in output_text:
                 output_text = output_text.split("</think", 1)[1].strip()
@@ -252,7 +230,7 @@ class SdxsPipeline(DiffusionPipeline):
         negative_prompt: Optional[Union[str, List[str]]] = None,
         height: int = 1152,
         width: int = 768,
-        num_inference_steps: int = 40,
+        num_inference_steps: int = 35,
         guidance_scale: float = 4.0,
         num_images_per_prompt: int = 1,
         seed: Optional[int] = None,
@@ -403,10 +381,10 @@ class SdxsPipeline(DiffusionPipeline):
         else:
             text_embeddings = prompt_embeds
 
-# 2. Prepare Timesteps (Bridging Diffusers Scheduler + Karras EDM)
+        # 2. Prepare Timesteps (Bridging Diffusers Scheduler + Karras EDM)
         # Ищем sigma_max/min в конфиге шедулера (если ты их туда добавил), 
         # если их там нет — используем константы Cosmos по умолчанию.
-        sigma_max = getattr(self.scheduler.config, "sigma_max", 80.0)
+        sigma_max = getattr(self.scheduler.config, "sigma_max", 10.0)
         sigma_min = getattr(self.scheduler.config, "sigma_min", 0.002)
         
         # Создаем наше любимое экспоненциальное убывание для лучшей детализации

transformer/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5b8df0f14160cc5614fa7eabb901b2c30b6566b04b940e5ea6858240bdb358b2
+oid sha256:07d6f1315358e1da06e51b1b023781bce56151d08169679d1212426ceb6e0b82
 size 3912877104