2512

sdxs_1b/diffusion_pytorch_model.safetensors

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

sdxs_flux/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:41dfe79753fadad9ba8e148ac73e1b61f23c7a1146cece5147c0041d613298bd
+oid sha256:5604baf577450654b1d025e8c31e476e42ebcb7d535eeb100917de6a31bf02fe
 size 3195253456

test.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7ee6048c4a59ddefaf190e5398f0dd7586746b0959db16971a118ca0a8e6e269
-size 8882423
+oid sha256:897429f78ba29492ec501a0b5b6224884068de6fc90a3976bc7b64f965b2b8e5
+size 9241758

train_flux.py

@@ -29,12 +29,12 @@ from transformers import AutoTokenizer, AutoModel
 # --------------------------- Параметры ---------------------------
 ds_path = "/workspace/sdxs/datasets/mjnj"
 project = "sdxs_flux"
-batch_size = 32
+batch_size = 256
 base_learning_rate = 4e-5 #2.7e-5
 min_learning_rate = 9e-6 #2.7e-5
-num_epochs = 10
-sample_interval_share = 10
-cfg_dropout = 0.9
+num_epochs = 20
+sample_interval_share = 3
+cfg_dropout = 0.5
 max_length = 192
 use_wandb = True
 use_comet_ml = False
@@ -96,8 +96,8 @@ lora_alpha = 64
 print("init")
 
 loss_ratios = {
-    "mse":   1.5,
-    "mae":   0.5,
+    "mse":   1.8,
+    "mae":   0.2,
 }
 median_coeff_steps = 256 
 
@@ -107,8 +107,6 @@ class MedianLossNormalizer:
         # нормируем доли на случай, если сумма != 1
         #s = sum(desired_ratios.values())
         #self.ratios = {k: (v / s) for k, v in desired_ratios.items()}
-        #self.buffers = {k: deque(maxlen=window_steps) for k in self.ratios.keys()}
-        #self.window = window_steps
         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
@@ -571,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
@@ -671,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
@@ -718,11 +718,7 @@ for epoch in range(start_epoch, start_epoch + num_epochs):
                 if accelerator.sync_gradients:
                     #with torch.amp.autocast('cuda', enabled=False):
                     grad_val = accelerator.clip_grad_norm_(unet.parameters(), clip_grad_norm)
-                    if grad_val is not None:
-                        grad = float(grad_val)
-                        print("grad_val is None")
-                    else:
-                        grad = 0.0 
+                    grad = float(grad_val)
                     optimizer.step()
                     lr_scheduler.step()
                     optimizer.zero_grad(set_to_none=True)