2511

train_vae_fdl.py

@@ -29,11 +29,11 @@ from collections import deque
 
 # --------------------------- Параметры ---------------------------
 ds_path            = "/workspace/d23"
-project            = "vae"
+project            = "vae4"
 batch_size         = 2
 base_learning_rate = 5e-5
 min_learning_rate  = 1e-5
-num_epochs         = 50
+num_epochs         = 25
 sample_interval_share = 5
 use_wandb          = True
 save_model         = True
@@ -41,8 +41,8 @@ use_decay          = True
 optimizer_type     = "adam8bit"
 dtype              = torch.float32
 
-model_resolution   = 256
-high_resolution    = 512 
+model_resolution   = 288
+high_resolution    = 576 
 limit              = 0
 save_barrier       = 1.3
 warmup_percent     = 0.005
@@ -51,9 +51,9 @@ beta2              = 0.997
 eps                = 1e-8
 clip_grad_norm     = 1.0
 mixed_precision    = "no"
-gradient_accumulation_steps = 2
+gradient_accumulation_steps = 4
 generated_folder   = "samples"
-save_as            = "vae4"
+save_as            = "vae5"
 num_workers        = 0
 device = None
 
@@ -67,10 +67,10 @@ kl_ratio           = 0.00
 # Доли лоссов
 loss_ratios = {
     "lpips": 0.70,#0.50,
-    "fdl" :  0.05,#0.25,   
+    "fdl" :  0.10,#0.25,   
     "edge":  0.05,
     "mse":   0.10,
-    "mae":   0.10,
+    "mae":   0.05,
     "kl":    0.00,  # активируем при full_training=True
 }
 median_coeff_steps = 1000
@@ -422,8 +422,8 @@ def generate_and_save_samples(step=None):
                 rec = temp_vae.decode(latents_mean).sample
 
         # Подгон размеров, если надо
-        if rec.shape[-2:] != orig_high.shape[-2:]:
-            rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
+        #if rec.shape[-2:] != orig_high.shape[-2:]:
+        #    rec = F.interpolate(rec, size=orig_high.shape[-2:], mode="bilinear", align_corners=False)
 
         # Сохраняем все real/decoded
         for i in range(rec.shape[0]):
@@ -437,8 +437,8 @@ def generate_and_save_samples(step=None):
         for i in range(rec.shape[0]):
             orig_full = orig_high[i:i+1].to(torch.float32)
             rec_full  = rec[i:i+1].to(torch.float32)
-            if rec_full.shape[-2:] != orig_full.shape[-2:]:
-                rec_full = F.interpolate(rec_full, size=orig_full.shape[-2:], mode="bilinear", align_corners=False)
+            #if rec_full.shape[-2:] != orig_full.shape[-2:]:
+            #    rec_full = F.interpolate(rec_full, size=orig_full.shape[-2:], mode="bilinear", align_corners=False)
             lpips_val = lpips_net(orig_full, rec_full).item()
             lpips_scores.append(lpips_val)
         avg_lpips = float(np.mean(lpips_scores))
@@ -497,8 +497,8 @@ for epoch in range(num_epochs):
                 latents = enc.latent_dist.mean if train_decoder_only else enc.latent_dist.sample()
                 rec = vae.decode(latents).sample
 
-            if rec.shape[-2:] != imgs.shape[-2:]:
-                rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
+            #if rec.shape[-2:] != imgs.shape[-2:]:
+            #    rec = F.interpolate(rec, size=imgs.shape[-2:], mode="bilinear", align_corners=False)
 
             rec_f32 = rec.to(torch.float32)
             imgs_f32 = imgs.to(torch.float32)

vae5/config.json

@@ -0,0 +1,48 @@
+{
+  "_class_name": "AsymmetricAutoencoderKL",
+  "_diffusers_version": "0.35.2",
+  "_name_or_path": "vae4",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": false,
+  "in_channels": 3,
+  "latent_channels": 16,
+  "layers_per_down_block": 2,
+  "layers_per_up_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 1.0,
+  "up_block_out_channels": [
+    128,
+    128,
+    256,
+    512,
+    512
+  ],
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

vae5/diffusion_pytorch_model.safetensors

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