2512

sdxs_1b/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:104b893ba36e1359f8468450653c0e00a7b2faeb1a42d3797747f0d0c06ccd0a
+oid sha256:3773b931314cc4356c69abfae833a9e948b3092b58d436a73d1cff8a36a74376
 size 4463672488

train.py

@@ -34,6 +34,7 @@ base_learning_rate = 4e-5 #2.7e-5
 min_learning_rate = 9e-6 #2.7e-5
 num_epochs = 10
 sample_interval_share = 20
+cfg_dropout = 0.5
 max_length = 192
 use_wandb = True
 use_comet_ml = False
@@ -95,8 +96,8 @@ lora_alpha = 64
 print("init")
 
 loss_ratios = {
-    "mse":   0.6,
-    "mae":   0.4,
+    "mse":   0.8,
+    "mae":   0.2,
 }
 median_coeff_steps = 256 
 
@@ -104,8 +105,9 @@ median_coeff_steps = 256
 class MedianLossNormalizer:
     def __init__(self, desired_ratios: dict, window_steps: int):
         # нормируем доли на случай, если сумма != 1
-        s = sum(desired_ratios.values())
-        self.ratios = {k: (v / s) for k, v in desired_ratios.items()}
+        #s = sum(desired_ratios.values())
+        #self.ratios = {k: (v / s) for k, v in desired_ratios.items()}
+        self.ratios = {k: float(v) for k, v in desired_ratios.items()}
         self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}
         self.window = window_steps
 
@@ -358,7 +360,7 @@ def collate_fn_simple(batch):
     raw_texts = [item["text"] for item in batch]
     texts = [
         "" if t.lower().startswith("zero") 
-        else "" if random.random() < 0.1
+        else "" if random.random() < cfg_dropout
         else t[1:].lstrip() if t.startswith(".")
         else t.replace("The image shows ", "").replace("The image is ", "").replace("This image captures ","").strip()
         for t in raw_texts
@@ -480,7 +482,7 @@ fixed_samples = get_fixed_samples_by_resolution(dataset)
 # --- [UPDATED] Функция для негативного эмбеддинга (возвращает 3 элемента) ---
 def get_negative_embedding(neg_prompt="", batch_size=1):
     if not neg_prompt:
-        hidden_dim = 2048 
+        hidden_dim = 1024 
         seq_len = max_length
         empty_emb = torch.zeros((batch_size, seq_len, hidden_dim), dtype=dtype, device=device)
         empty_mask = torch.ones((batch_size, seq_len), dtype=torch.int64, device=device)
@@ -567,6 +569,8 @@ def generate_and_save_samples(fixed_samples_cpu, uncond_data, step):
                 latents = scheduler.step(flow, t, latents).prev_sample
         
             current_latents = latents
+            if step==0:
+                current_latents = sample_latents
 
             latent_for_vae = current_latents.detach() / scaling_factor + shift_factor
             decoded = vae.decode(latent_for_vae.to(torch.float32)).sample
@@ -667,9 +671,9 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
             # шум
             noise = torch.randn_like(latents, dtype=latents.dtype)
             # берём t из [0, 1]
-            #t = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
-            t = torch.sigmoid(torch.randn_like(u))
+            t = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
+            #u = torch.rand(latents.shape[0], device=latents.device, dtype=latents.dtype)
+            #t = torch.sigmoid(torch.randn_like(u))
             
             # интерполяция между x0 и шумом
             noisy_latents = (1.0 - t.view(-1, 1, 1, 1)) * latents + t.view(-1, 1, 1, 1) * noise