Spaces:

Rohan3
/

flickr8k-backend

Sleeping

App Files Files Community

Rohan3 commited on Apr 21

Commit

a625e96

1 Parent(s): 87b5061

Updated: VAE, UNet, config, text embeddings, model and main

Browse files

Files changed (10) hide show

app/main.py +4 -77
app/model.py +3 -5
core/config.py +11 -7
core/dataloader.py +10 -9
core/sample_ddim.py +2 -2
core/text_embeddings.py +4 -6
core/train_unet.py +17 -24
core/train_vae.py +25 -27
core/unet.py +19 -35
core/vae.py +32 -16

app/main.py CHANGED Viewed

@@ -36,9 +36,10 @@ app.add_middleware(
 class GenerateRequest(BaseModel):
     caption: str = Field(..., example="a white dog running in snow")
     num_images: int = Field(4, ge=1, le=8)
-    num_steps: int = Field(50, ge=10, le=200)
-    guidance_scale: float = Field(7.5, ge=1.0, le=20.0)
     seed: int = Field(42)
 def tensor_to_pil(img_tensor: torch.Tensor) -> Image.Image:
     img = img_tensor.clamp(0, 1)
@@ -78,83 +79,9 @@ async def generate(req: GenerateRequest, _=Depends(verify_key)):
             num_steps=req.num_steps,
             guidance_scale=req.guidance_scale,
             seed=req.seed,
         )
         b64_images = [tensor_to_base64(img) for img in images]
         return GenerateResponse(images=b64_images, num_generated=len(b64_images))
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-# ── Single image (num_images=1) ──────────────────────────────────────
-@app.post("/generate/image", response_class=StreamingResponse)
-async def generate_image(req: GenerateRequest):
-    try:
-        req.num_images = 1
-        images = pipeline.generate(
-            caption=req.caption,
-            num_images=1,
-            num_steps=req.num_steps,
-            guidance_scale=req.guidance_scale,
-            seed=req.seed,
-        )
-        buf = io.BytesIO()
-        tensor_to_pil(images[0]).save(buf, format="PNG")
-        buf.seek(0)
-        return StreamingResponse(buf, media_type="image/png")
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-# ── Multiple images as ZIP ───────────────────────────────────────────
-@app.post("/generate/zip", response_class=StreamingResponse)
-async def generate_zip(req: GenerateRequest):
-    try:
-        images = pipeline.generate(
-            caption=req.caption,
-            num_images=req.num_images,
-            num_steps=req.num_steps,
-            guidance_scale=req.guidance_scale,
-            seed=req.seed,
-        )
-        buf = io.BytesIO()
-        with zipfile.ZipFile(buf, "w") as zf:
-            for i, img_tensor in enumerate(images):
-                img_buf = io.BytesIO()
-                tensor_to_pil(img_tensor).save(img_buf, format="PNG")
-                zf.writestr(f"image_{i}.png", img_buf.getvalue())
-        buf.seek(0)
-        return StreamingResponse(buf, media_type="application/zip",
-                                 headers={"Content-Disposition": "attachment; filename=images.zip"})
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-# ── Preview page: see all images directly in browser ─────────────────
-@app.post("/generate/preview", response_class=HTMLResponse)
-async def generate_preview(req: GenerateRequest):
-    try:
-        images = pipeline.generate(
-            caption=req.caption,
-            num_images=req.num_images,
-            num_steps=req.num_steps,
-            guidance_scale=req.guidance_scale,
-            seed=req.seed,
-        )
-        img_tags = ""
-        for img_tensor in images:
-            transform = T.ToPILImage()
-            img = transform(img_tensor)
-            img.show()
-            buf = io.BytesIO()
-            tensor_to_pil(img_tensor).save(buf, format="PNG")
-            b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
-            img_tags += f'<img src="data:image/png;base64,{b64}" style="width:256px; margin:5px;"/>'
-        html = f"""
-        <html>
-            <body style="background:#111; display:flex; flex-wrap:wrap;">
-                <h2 style="color:white; width:100%;">Caption: {req.caption}</h2>
-                {img_tags}
-            </body>
-        </html>
-        """
-        return HTMLResponse(content=html)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

 class GenerateRequest(BaseModel):
     caption: str = Field(..., example="a white dog running in snow")
     num_images: int = Field(4, ge=1, le=8)
+    num_steps: int = Field(30, ge=10, le=100)
+    guidance_scale: float = Field(5, ge=1.0, le=20.0)
     seed: int = Field(42)
+    eta: float = Field(0, ge=0.0, le=1.0)
 def tensor_to_pil(img_tensor: torch.Tensor) -> Image.Image:
     img = img_tensor.clamp(0, 1)
             num_steps=req.num_steps,
             guidance_scale=req.guidance_scale,
             seed=req.seed,
+            eta=req.eta,
         )
         b64_images = [tensor_to_base64(img) for img in images]
         return GenerateResponse(images=b64_images, num_generated=len(b64_images))
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

app/model.py CHANGED Viewed

@@ -84,14 +84,12 @@ class LDMPipeline:
         per_token_contextual = self.text_model.ln_final(x) # (B, T, D) = (1, 77, 1024)
         return per_token_contextual.squeeze(0) # (77, 1024)`
-    def generate(self, caption: str, num_images: int = 4,
-                 num_steps: int = 50, guidance_scale: float = 7.5,
-                 seed: int = 42):
         seed_everything(seed)
         caption = caption.strip()
         if caption.endswith("."):
             caption = caption.rstrip(".")
-        caption = caption.lower()
         embedding = self.get_text_embedding(caption).unsqueeze(0)
         latents = ddim_sample(
             unet=self.unet,
@@ -101,7 +99,7 @@ class LDMPipeline:
             embedding=embedding,
             guidance_scale=guidance_scale,
             num_steps=num_steps,
-            eta=0.0,
             device=self.device
         )
         latents = latents * latent_std

         per_token_contextual = self.text_model.ln_final(x) # (B, T, D) = (1, 77, 1024)
         return per_token_contextual.squeeze(0) # (77, 1024)`
+    def generate(self, caption: str, num_images: int = 4, num_steps: int = 50, guidance_scale: float = 7.5, seed: int = 42, eta: float = 0):
         seed_everything(seed)
         caption = caption.strip()
         if caption.endswith("."):
             caption = caption.rstrip(".")
+        # caption = caption.lower()
         embedding = self.get_text_embedding(caption).unsqueeze(0)
         latents = ddim_sample(
             unet=self.unet,
             embedding=embedding,
             guidance_scale=guidance_scale,
             num_steps=num_steps,
+            eta=eta,
             device=self.device
         )
         latents = latents * latent_std

core/config.py CHANGED Viewed

@@ -1,26 +1,29 @@
 tqdm_colors = ["#0867C5", "green", "blue", "red", "yellow", "#6707B0", "#FFDD22", "#00FFEF", "#DC442C"]
 data_dir = "./backend/core/data"
-image_res = 128
 image_dir = f"{data_dir}/Images"
 resized_img_dir = f"{data_dir}/ResizedImages_{image_res}"
-vae_batch_size = 8
 vae_group_size = 4
 vae_num_epochs = 1000
 vae_stopping_patience = 30
 vae_latent_channels = 4 # 16, 8
 vae_latent_dim = 32
 vae_beta_kld = 1e-3
-vae_optim_lr = 1e-4
 vae_lambda_tvl = 0 # 1e-3
 vae_lpips_weight = 1e-1
 vae_dropout = 0.
 vae_checkpoint_dir = "./backend/core/checkpoints/vae"
-vae_weight = f"attn_test_best_{image_res}_3_4_32_128_{vae_latent_channels}_beta_{vae_beta_kld}_tvl_{vae_lambda_tvl}_batch_{vae_batch_size}_lr_{vae_optim_lr}.pth"
 latent_dir = f"./backend/core/latents_{image_res}"
 latent_scaled_dir = f"./backend/core/latents_scaled_{image_res}"
 latent_norm_dir = f"./backend/core/latents_norm_{image_res}"
 latent_recon_images = f"./backend/core/recon_img_{vae_latent_channels}_{vae_latent_dim}_{vae_latent_dim}_res_{image_res}"
-latent_mu = 0.021192772313952446; latent_std = 0.9767765402793884 # Latents
 # Mean over latent: 0.021192772313952446
 # STD over latent: 0.9767765402793884
 # Latent Scale: 1.023775577545166
@@ -39,19 +42,20 @@ text_captions_dir = f"{data_dir}/captions.txt"
 unet_batch_size = 32
 embedding_dim = 1024
 embedding_dir = f"./backend/core/embeddings_77_{embedding_dim}"
 embedding_model = "ViT-g-14" # or "ViT-B-16"
 embedding_pretrained = "laion2b_s12b_b42k" # or "openai"
 unet_pred_type = "v_prediction" # v_prediction or epsilon
 unet_checkpoint_dir = f"./backend/core/checkpoints/ldm"
 unet_max_steps = 256_000
 # new_lr = old_lr * (batch_size_new / batch_size_old)
-unet_optim_lr = 1e-4
 unet_group_size = 32
 unet_beta_schedule = "squaredcos_cap_v2" # "linear" or "squaredcos_cap_v2"
 unet_dropout = 0.
 attn_dropout = 0.
 unet_train_timesteps = 1000
-ddim_guidace_scale = 3
 ddim_num_sampling_steps = 100
 ddim_img_dir = "./backend/core/ddim_recon_img"
 unet_val_embeddings_dir = "./backend/core/embeddings_val"

 tqdm_colors = ["#0867C5", "green", "blue", "red", "yellow", "#6707B0", "#FFDD22", "#00FFEF", "#DC442C"]
 data_dir = "./backend/core/data"
+image_res = 256
 image_dir = f"{data_dir}/Images"
 resized_img_dir = f"{data_dir}/ResizedImages_{image_res}"
+vae_batch_size = 4
 vae_group_size = 4
 vae_num_epochs = 1000
 vae_stopping_patience = 30
 vae_latent_channels = 4 # 16, 8
 vae_latent_dim = 32
 vae_beta_kld = 1e-3
+vae_optim_lr = 5e-5
 vae_lambda_tvl = 0 # 1e-3
 vae_lpips_weight = 1e-1
 vae_dropout = 0.
 vae_checkpoint_dir = "./backend/core/checkpoints/vae"
+# vae_weight = f"attn_test_best_{image_res}_3_32_128_256_512_{vae_latent_channels}_beta_{vae_beta_kld}_tvl_{vae_lambda_tvl}_batch_{vae_batch_size}_lr_{vae_optim_lr}.pth"
+vae_weight = f"ssim_lpips_attn_test_best_{image_res}_3_32_128_256_512_{vae_latent_channels}_beta_{vae_beta_kld}_tvl_{vae_lambda_tvl}_batch_{vae_batch_size}_lr_{vae_optim_lr}.pth"
 latent_dir = f"./backend/core/latents_{image_res}"
 latent_scaled_dir = f"./backend/core/latents_scaled_{image_res}"
 latent_norm_dir = f"./backend/core/latents_norm_{image_res}"
 latent_recon_images = f"./backend/core/recon_img_{vae_latent_channels}_{vae_latent_dim}_{vae_latent_dim}_res_{image_res}"
+# latent_mu = 0.021192772313952446; latent_std = 0.9767765402793884 # Latents
+# latent_mu = 0.0007418220047838986; latent_std = 0.9822604060173035 # Latents scale = 1.0180599689483643
+latent_mu =  -0.08573896437883377; latent_std = 1.2452856302261353 # Latents scale = 0.8030286431312561
 # Mean over latent: 0.021192772313952446
 # STD over latent: 0.9767765402793884
 # Latent Scale: 1.023775577545166
 unet_batch_size = 32
 embedding_dim = 1024
 embedding_dir = f"./backend/core/embeddings_77_{embedding_dim}"
+null_embedding_dir = f"./backend/core/embeddings_77_{embedding_dim}/null_embedding.pt"
 embedding_model = "ViT-g-14" # or "ViT-B-16"
 embedding_pretrained = "laion2b_s12b_b42k" # or "openai"
 unet_pred_type = "v_prediction" # v_prediction or epsilon
 unet_checkpoint_dir = f"./backend/core/checkpoints/ldm"
 unet_max_steps = 256_000
 # new_lr = old_lr * (batch_size_new / batch_size_old)
+unet_optim_lr = 1e-4 # Changed to 5e-5 after 1103080 steps
 unet_group_size = 32
 unet_beta_schedule = "squaredcos_cap_v2" # "linear" or "squaredcos_cap_v2"
 unet_dropout = 0.
 attn_dropout = 0.
 unet_train_timesteps = 1000
+ddim_guidace_scale = 8
 ddim_num_sampling_steps = 100
 ddim_img_dir = "./backend/core/ddim_recon_img"
 unet_val_embeddings_dir = "./backend/core/embeddings_val"

core/dataloader.py CHANGED Viewed

@@ -21,13 +21,13 @@ class ImageDataset(Dataset):
 def image_dataloader():
     image_paths = [os.path.join(resized_img_dir, path) for path in os.listdir(resized_img_dir)[:] if path.endswith(".jpg")]
     dataset = ImageDataset(image_paths)
-    g = torch.Generator()
-    g.manual_seed(42)
     train_size = int(0.8 * len(dataset))
     test_size = len(dataset) - train_size
     train_set, test_set = random_split(dataset, [train_size, test_size])
-    train_loader = DataLoader(train_set, batch_size=vae_batch_size, shuffle=True, generator=g, worker_init_fn=seed_worker, num_workers=min(8, os.cpu_count()), pin_memory=True, persistent_workers=True, prefetch_factor=16)
-    test_loader = DataLoader(test_set, batch_size=vae_batch_size, shuffle=False, generator=g, worker_init_fn=seed_worker, num_workers=min(8, os.cpu_count()), pin_memory=True, persistent_workers=True, prefetch_factor=16)
     return train_loader, test_loader
 class LatentEmbeddingsDataset(Dataset):
@@ -39,20 +39,21 @@ class LatentEmbeddingsDataset(Dataset):
     def __getitem__(self, index):
         latent_index = index // self.num_variants
         variant_index = index % self.num_variants
         file_path = self.file_names[latent_index]
         latent_data = torch.load(f"{latent_scaled_dir}/{file_path}", map_location="cpu", weights_only=True)
         embedding_data = torch.load(f"{embedding_dir}/{file_path[:-3]}_{variant_index}.pt", map_location="cpu", weights_only=True)
         assert embedding_data.shape == torch.Size([77, 1024]), f"Unexpected embedding shape: {embedding_data.shape}"
         return (latent_data, embedding_data)
 def latent_embedding_dataloader():
-    file_names = [path for path in os.listdir(latent_scaled_dir) if path.endswith(".pt")][:]
     dataset = LatentEmbeddingsDataset(file_names)
-    g = torch.Generator()
-    g.manual_seed(42)
-    train_loader = DataLoader(dataset, batch_size=unet_batch_size, shuffle=True, generator=g, worker_init_fn=seed_worker, num_workers=min(8, os.cpu_count()), pin_memory=True, persistent_workers=True, prefetch_factor=32)
     return train_loader
 def seed_worker(worker_id):
-    worker_seed = 42 + worker_id
     np.random.seed(worker_seed)
     random.seed(worker_seed)

 def image_dataloader():
     image_paths = [os.path.join(resized_img_dir, path) for path in os.listdir(resized_img_dir)[:] if path.endswith(".jpg")]
     dataset = ImageDataset(image_paths)
+    # g = torch.Generator()
+    # g.manual_seed(42)
     train_size = int(0.8 * len(dataset))
     test_size = len(dataset) - train_size
     train_set, test_set = random_split(dataset, [train_size, test_size])
+    train_loader = DataLoader(train_set, batch_size=vae_batch_size, shuffle=True, worker_init_fn=seed_worker, num_workers=min(8, os.cpu_count()), pin_memory=True, persistent_workers=True, prefetch_factor=4)
+    test_loader = DataLoader(test_set, batch_size=vae_batch_size, shuffle=False, worker_init_fn=seed_worker, num_workers=min(8, os.cpu_count()), pin_memory=True, persistent_workers=True, prefetch_factor=4)
     return train_loader, test_loader
 class LatentEmbeddingsDataset(Dataset):
     def __getitem__(self, index):
         latent_index = index // self.num_variants
         variant_index = index % self.num_variants
+        # variant_index = random.randint(0, self.num_variants - 1)
         file_path = self.file_names[latent_index]
         latent_data = torch.load(f"{latent_scaled_dir}/{file_path}", map_location="cpu", weights_only=True)
         embedding_data = torch.load(f"{embedding_dir}/{file_path[:-3]}_{variant_index}.pt", map_location="cpu", weights_only=True)
         assert embedding_data.shape == torch.Size([77, 1024]), f"Unexpected embedding shape: {embedding_data.shape}"
         return (latent_data, embedding_data)
 def latent_embedding_dataloader():
+    file_names = sorted(path for path in os.listdir(latent_scaled_dir) if path.endswith(".pt"))[:]
+    # file_names = [path for path in os.listdir(latent_scaled_dir) if path.endswith(".pt")][:16]
     dataset = LatentEmbeddingsDataset(file_names)
+    train_loader = DataLoader(dataset, batch_size=unet_batch_size, shuffle=True, worker_init_fn=seed_worker, num_workers=min(8, os.cpu_count()), pin_memory=True, persistent_workers=True, prefetch_factor=32)
     return train_loader
 def seed_worker(worker_id):
+    # worker_seed = 42 + worker_id
+    worker_seed = torch.initial_seed() % 2**32
     np.random.seed(worker_seed)
     random.seed(worker_seed)

core/sample_ddim.py CHANGED Viewed

@@ -17,7 +17,7 @@ def ddim_sample(unet, noise_scheduler, shape, null_embedding=None, x_start=None,
     # x = torch.randn(shape, device=device) # start from pure noise
     noise_scheduler.set_timesteps(num_steps) # set DDIM timesteps
     embedding = embedding.expand(x.shape[0], -1, -1) if embedding != None else None
-    if null_embedding == None: null_embedding = torch.load(f"{embedding_dir}/null_embedding.pt", map_location=device, weights_only=True).unsqueeze(0) # [1, 77, 1024]
     null_embedding = null_embedding.expand(x.shape[0], -1, -1)
     for t in tqdm(noise_scheduler.timesteps, desc=f"Sampling timesteps: ", colour=tqdm_colors[-1]):
         t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
@@ -32,7 +32,7 @@ def ddim_sample(unet, noise_scheduler, shape, null_embedding=None, x_start=None,
             target_std = max(std_uncond, std_cond) # Target std is usually the max of the two individual predictions or just the conditional std.
             factor = target_std / std_cfg # Factor to bring the CFG output back into range
             model_out_rescaled = model_out * factor
-            phi = 0.7 # Final blend (phi=0.7 is standard)
             model_out = phi * model_out_rescaled + (1 - phi) * model_out
         else: model_out = unet(x, t_batch, embedding)
         x = noise_scheduler.step(model_output=model_out, timestep=t, sample=x, eta=eta).prev_sample # DDIM step

     # x = torch.randn(shape, device=device) # start from pure noise
     noise_scheduler.set_timesteps(num_steps) # set DDIM timesteps
     embedding = embedding.expand(x.shape[0], -1, -1) if embedding != None else None
+    if null_embedding == None: null_embedding = torch.load(null_embedding_dir, map_location=device, weights_only=True).unsqueeze(0) # [1, 77, 1024]
     null_embedding = null_embedding.expand(x.shape[0], -1, -1)
     for t in tqdm(noise_scheduler.timesteps, desc=f"Sampling timesteps: ", colour=tqdm_colors[-1]):
         t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
             target_std = max(std_uncond, std_cond) # Target std is usually the max of the two individual predictions or just the conditional std.
             factor = target_std / std_cfg # Factor to bring the CFG output back into range
             model_out_rescaled = model_out * factor
+            phi = 0.7 # Final blend
             model_out = phi * model_out_rescaled + (1 - phi) * model_out
         else: model_out = unet(x, t_batch, embedding)
         x = noise_scheduler.step(model_output=model_out, timestep=t, sample=x, eta=eta).prev_sample # DDIM step

core/text_embeddings.py CHANGED Viewed

@@ -11,10 +11,8 @@ def save_embeddings(model, tokenizer, device):
         last_img_name = ""; counter = 0
         for l in tqdm(f.readlines(), desc=f"Progress"):
             img_name, caption = l.strip().split(".jpg,")
-            if caption[-1]==".":
-                caption = caption.strip().rstrip('.')
-                # caption = caption[:-2]
-            caption = caption.lower()
             if len(last_img_name) == 0 or last_img_name != img_name: last_img_name = img_name; counter = 0
             else: counter += 1
             img_name += f"_{counter}"
@@ -24,7 +22,7 @@ def save_embeddings(model, tokenizer, device):
 def save_null_embedding(model, tokenizer, device):
     embedding = get_text_embedding(model, tokenizer, "", device)
-    torch.save(embedding.cpu(), f"./{embedding_dir}/null_embedding.pt")
     print("Saved null embedding")
 def save_val_embeddings(model, tokenizer, device):
@@ -73,7 +71,7 @@ if __name__ == "__main__":
     save_val_embeddings(model, tokenizer, device)
 # Check if embeddings are equal
-# emb1 = torch.load("./backend/best_checkpoints/null_embedding/null_embedding.pt", map_location="cpu", weights_only=True)
 # emb2 = torch.load("./backend/core/embeddings_77_1024/null_embedding.pt", map_location="cpu", weights_only=True)
 # print("Max diff:", (emb1 - emb2).abs().max().item())
 # print("Are equal:", torch.allclose(emb2, emb1, atol=1e-5))

         last_img_name = ""; counter = 0
         for l in tqdm(f.readlines(), desc=f"Progress"):
             img_name, caption = l.strip().split(".jpg,")
+            caption = caption.strip()
+            if caption[-1]==".": caption = caption.rstrip('.')
             if len(last_img_name) == 0 or last_img_name != img_name: last_img_name = img_name; counter = 0
             else: counter += 1
             img_name += f"_{counter}"
 def save_null_embedding(model, tokenizer, device):
     embedding = get_text_embedding(model, tokenizer, "", device)
+    torch.save(embedding.cpu(), null_embedding_dir)
     print("Saved null embedding")
 def save_val_embeddings(model, tokenizer, device):
     save_val_embeddings(model, tokenizer, device)
 # Check if embeddings are equal
+# emb1 = torch.load("./backend/core/null_embedding.pt", map_location="cpu", weights_only=True)
 # emb2 = torch.load("./backend/core/embeddings_77_1024/null_embedding.pt", map_location="cpu", weights_only=True)
 # print("Max diff:", (emb1 - emb2).abs().max().item())
 # print("Are equal:", torch.allclose(emb2, emb1, atol=1e-5))

core/train_unet.py CHANGED Viewed

@@ -16,7 +16,7 @@ from sample_ddim import gen_n_sampled_img, gen_val_sampled_img
 import numpy as np
 import random
-def train_unet_ddpm_simple(use_checkpoint=False):
     os.makedirs(unet_checkpoint_dir, exist_ok=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print("DEVICE", device)
@@ -35,22 +35,19 @@ def train_unet_ddpm_simple(use_checkpoint=False):
     save_every = epoch_size * 8
     val_every = epoch_size * 4
     ema_warmup_steps = epoch_size * 20
-    null_embedding = (torch.load(f"{embedding_dir}/null_embedding.pt", map_location=device, weights_only=True).unsqueeze(0))
     ema = EMA(unet, decay=0.999, warmup_steps=ema_warmup_steps)
     scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
     global_step = loaded_step = 0
     total_mse = total_samples = 0
-    checkpoint_path = "./backend/core/checkpoints/ldm/ema_step_586960_0.298282.pth"
-    if use_checkpoint and os.path.exists(checkpoint_path):
-        loaded_step = 586960
         print("Loaded checkpoint successfully")
         checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
         if "rng_state" in checkpoint: torch.set_rng_state(checkpoint["rng_state"].cpu())
         if torch.cuda.is_available() and "cuda_rng_state" in checkpoint:
             torch.cuda.set_rng_state(checkpoint["cuda_rng_state"].cpu())
-            # torch.cuda.set_rng_state_all(checkpoint["cuda_rng_state"].cpu())
-            # np.random.set_state(checkpoint["cuda_rng_state"].cpu())
-            # random.setstate(checkpoint["cuda_rng_state"].cpu())
         unet.load_state_dict(checkpoint["unet"], strict=True)
         global_step = checkpoint["step"]
         optimizer.load_state_dict(checkpoint["optimizer"])
@@ -58,17 +55,19 @@ def train_unet_ddpm_simple(use_checkpoint=False):
         noise_scheduler = DDPMScheduler.from_pretrained("./checkpoints/noise_scheduler")
         ema.shadow = {k: v.to(device) for k, v in checkpoint["ema"].items()}
         ema.num_updates = checkpoint["ema_num_updates"]
     unet.train()
     while global_step < max_steps:
         for latent, embedding in tqdm(train_loader, desc="Train: ", colour=tqdm_colors[5]):
             if global_step >= max_steps: break
             latent, embedding = latent.to(device), embedding.to(device)
             assert null_embedding.shape[1:] == embedding.shape[1:], f"{null_embedding.shape}"
             batch_size = latent.size(0)
             t = torch.randint(0, noise_scheduler.config.num_train_timesteps, (batch_size,), device=device).long()
             noise = torch.randn_like(latent)
             noised_latent = noise_scheduler.add_noise(latent, noise, t)
-            mask = torch.rand(batch_size, 1, 1, device=device) < 0.2
             embedding = torch.where(mask, null_embedding.expand(batch_size, -1, -1), embedding)
             optimizer.zero_grad(set_to_none=True)
             with torch.amp.autocast(device_type="cuda", enabled=torch.cuda.is_available()):
@@ -83,20 +82,14 @@ def train_unet_ddpm_simple(use_checkpoint=False):
             scaler.update()
             if scaler.get_scale() >= old_scale:
                 ema.update(unet)
                 global_step += 1
-            # scaler.scale(loss).backward()
-            # scaler.unscale_(optimizer)
-            # torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)
-            # scaler.step(optimizer)
-            # scaler.update()
-            # # EMA update
-            # ema.update(unet)
-            # global_step += 1
-            total_mse += loss.item() * batch_size
-            total_samples += batch_size
-            if global_step % log_every == 0  and (global_step != loaded_step):
-                print(f" Step {global_step} | Loss: {total_mse / total_samples:.6f}")
-                if global_step % save_every == 0  and (global_step != loaded_step):
                     torch.save({
                         "step": global_step,
                         "unet": unet.state_dict(),
@@ -109,7 +102,7 @@ def train_unet_ddpm_simple(use_checkpoint=False):
                     }, f"{unet_checkpoint_dir}/ema_step_{global_step}_{total_mse / total_samples:.6f}.pth")
                     noise_scheduler.save_pretrained(f"./checkpoints/noise_scheduler")
                 total_samples = total_mse = 0
-            if global_step % val_every == 0  and (global_step != loaded_step):
                 unet.eval()
                 applied_ema = False
                 with torch.no_grad():
@@ -126,4 +119,4 @@ def train_unet_ddpm_simple(use_checkpoint=False):
 if __name__ == "__main__":
     seed_everything(42)
-    train_unet_ddpm_simple(use_checkpoint=False)

 import numpy as np
 import random
+def train_unet_ddpm_simple(use_checkpoint=False, checkpoint_path=None):
     os.makedirs(unet_checkpoint_dir, exist_ok=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print("DEVICE", device)
     save_every = epoch_size * 8
     val_every = epoch_size * 4
     ema_warmup_steps = epoch_size * 20
+    null_embedding = (torch.load(null_embedding_dir, map_location=device, weights_only=True).unsqueeze(0))
     ema = EMA(unet, decay=0.999, warmup_steps=ema_warmup_steps)
     scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
+    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_steps, eta_min=1e-6)
     global_step = loaded_step = 0
     total_mse = total_samples = 0
+    if use_checkpoint and checkpoint_path and os.path.exists(checkpoint_path):
+        loaded_step = int(checkpoint_path.split("step_")[1].split("_")[0])
         print("Loaded checkpoint successfully")
         checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
         if "rng_state" in checkpoint: torch.set_rng_state(checkpoint["rng_state"].cpu())
         if torch.cuda.is_available() and "cuda_rng_state" in checkpoint:
             torch.cuda.set_rng_state(checkpoint["cuda_rng_state"].cpu())
         unet.load_state_dict(checkpoint["unet"], strict=True)
         global_step = checkpoint["step"]
         optimizer.load_state_dict(checkpoint["optimizer"])
         noise_scheduler = DDPMScheduler.from_pretrained("./checkpoints/noise_scheduler")
         ema.shadow = {k: v.to(device) for k, v in checkpoint["ema"].items()}
         ema.num_updates = checkpoint["ema_num_updates"]
+        # for _ in range(global_step): scheduler.step()
     unet.train()
     while global_step < max_steps:
         for latent, embedding in tqdm(train_loader, desc="Train: ", colour=tqdm_colors[5]):
             if global_step >= max_steps: break
             latent, embedding = latent.to(device), embedding.to(device)
             assert null_embedding.shape[1:] == embedding.shape[1:], f"{null_embedding.shape}"
+            assert embedding.dim() == 3, f"Expected 3D embedding (B, seq, dim), got {embedding.shape}"
             batch_size = latent.size(0)
             t = torch.randint(0, noise_scheduler.config.num_train_timesteps, (batch_size,), device=device).long()
             noise = torch.randn_like(latent)
             noised_latent = noise_scheduler.add_noise(latent, noise, t)
+            mask = torch.rand(batch_size, 1, 1, device=device) < 0.1
             embedding = torch.where(mask, null_embedding.expand(batch_size, -1, -1), embedding)
             optimizer.zero_grad(set_to_none=True)
             with torch.amp.autocast(device_type="cuda", enabled=torch.cuda.is_available()):
             scaler.update()
             if scaler.get_scale() >= old_scale:
                 ema.update(unet)
+                # scheduler.step()
                 global_step += 1
+                total_mse += loss.item() * batch_size
+                total_samples += batch_size
+            if global_step % log_every == 0 and (global_step != loaded_step):
+                # print(f" Step {global_step} | Loss: {total_mse / total_samples:.6f} | LR: {scheduler.get_last_lr()[0]:.2e}")
+                if total_samples > 0: print(f" Step {global_step} | Loss: {total_mse / total_samples:.6f}")
+                if global_step % save_every == 0 and (global_step != loaded_step):
                     torch.save({
                         "step": global_step,
                         "unet": unet.state_dict(),
                     }, f"{unet_checkpoint_dir}/ema_step_{global_step}_{total_mse / total_samples:.6f}.pth")
                     noise_scheduler.save_pretrained(f"./checkpoints/noise_scheduler")
                 total_samples = total_mse = 0
+            if global_step % val_every == 0 and (global_step != loaded_step):
                 unet.eval()
                 applied_ema = False
                 with torch.no_grad():
 if __name__ == "__main__":
     seed_everything(42)
+    train_unet_ddpm_simple(use_checkpoint=True, checkpoint_path="./backend\core\checkpoints\ldm\ema_step_1163800_0.262618.pth")

core/train_vae.py CHANGED Viewed

@@ -1,5 +1,6 @@
 import sys, os
 sys.path.insert(0, os.path.dirname(__file__))
 import torch
 from torch import nn, optim
 import torch.nn.functional as F
@@ -14,12 +15,6 @@ from sample_vae import reconstruct
 import lpips
 from pytorch_msssim import SSIM
-# def vae_loss(recon_x, x, mu, logvar, beta_kld):
-#     # loss = F.l1_loss(recon_x, x, reduction="sum") / recon_x.size(0)
-#     loss = F.mse_loss(recon_x, x, reduction="sum") / recon_x.size(0)
-#     kld = torch.mean(-0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1))
-#     tvl = total_variance_loss(recon_x)
-#     return loss + kld*beta_kld + tvl*vae_lambda_tvl, loss, kld, tvl
 # # Simple PatchGAN Discriminator
 # class PatchDiscriminator(nn.Module):
@@ -73,22 +68,25 @@ def total_variance_loss(x):
     tvl_w = torch.pow(x[:, :, :, 1:] - x[:, :, :, :-1], 2).sum() # TV-L2
     return (tvl_h + tvl_w) / x.size(0)
-def get_annealed_beta(epoch, warmup_epochs=100, max_beta=1.0): return max_beta * min(vae_beta_kld, epoch / warmup_epochs)
 # def vae_loss(recon_x, x, mu, logvar, lpips_fn, ssim_fn, beta_kld=1.0):
-def vae_loss(recon_x, x, mu, logvar, beta_kld=1.0):
     # mse = torch.mean((recon_x - x) ** 2)
     # kld_loss = torch.mean(-0.5 * (1 + logvar - mu.pow(2) - logvar.exp()))
     # b, c, h, w = x.shape
-    # tvl = total_variance_loss(recon_x) if vae_lambda_tvl > 0 else 0.0
     mse = F.mse_loss(recon_x, x, reduction="sum") / recon_x.size(0)
     kld_loss = torch.sum(-0.5 * (1 + logvar - mu.pow(2) - logvar.exp())) / recon_x.size(0)
-    # ssim_loss = 1 - ssim_fn(recon_x, x)
-    # lpips_loss = lpips_fn(recon_x * 2 - 1, x * 2 - 1).mean()
-    # total_loss = mse + (beta_kld * kld_loss) + (vae_lambda_tvl * tvl) + (vae_lpips_weight * lpips_loss) + ssim_loss * 0.1
-    # return total_loss, mse, kld_loss, tvl, lpips_loss, ssim_loss
-    total_loss = mse + (beta_kld * kld_loss)
-    return total_loss, mse, kld_loss, 0, 0, 0
 def train_test_vae():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -98,14 +96,14 @@ def train_test_vae():
     scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5)
     os.makedirs(f"{vae_checkpoint_dir}", exist_ok=True)
     early_stopping_counter = 0
-    # global_step = 0
-    # warmup_steps = int(0.1 * len(train_loader) * vae_num_epochs)
     best_test_loss = float("inf")
     train_bce_loss = []
     train_kld_loss = []
-    # lpips_fn = lpips.LPIPS(net='vgg').to(device)
-    # lpips_fn.eval()
-    # ssim_fn = SSIM(data_range=1, size_average=True, channel=3)
     scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
     # discriminator = PatchDiscriminator().to(device)
     # optimizer_D = optim.AdamW(discriminator.parameters(), lr=vae_optim_lr, betas=(0.5, 0.999))
@@ -126,18 +124,19 @@ def train_test_vae():
                 # loss, rec, kld, tvl, lpips_loss, gan_loss = vae_loss(recon_x, x, mu, logvar, discriminator, lpips_fn, beta_kld=get_annealed_beta(epoch)) # Update VAE (generator)
                 # loss, rec, kld, tvl, lpips_loss = vae_loss(recon_x, x, mu, logvar, lpips_fn, beta_kld=min(1, global_step / warmup_steps)) # Update VAE (generator)
                 # loss, rec, kld, tvl, lpips_loss, ssim_loss = vae_loss(recon_x, x, mu, logvar, lpips_fn, ssim_fn, beta_kld=get_annealed_beta(epoch)) # Update VAE (generator)
-                loss, rec, kld, tvl, lpips_loss, ssim_loss = vae_loss(recon_x, x, mu, logvar, beta_kld=get_annealed_beta(epoch)) # Update VAE (generator)
             optimizer.zero_grad(set_to_none=True)
             scaler.scale(loss).backward()
             scaler.step(optimizer)
             scaler.update()
             train_loss += loss.item()
             total_rec += rec.item()
             total_kld += kld.item()
-            # total_tvl += tvl
-            # total_lpips += lpips_loss
-            # total_lpips += lpips_loss.item()
-            # total_ssim += ssim_loss.item()
             # total_gan += gan_loss.item()
         avg_train_loss = train_loss / len(train_loader)
         avg_rec = total_rec / len(train_loader)
@@ -157,7 +156,7 @@ def train_test_vae():
                     recon_x, mu, logvar = vae(x)
                     # loss, *_ = vae_loss(recon_x, x, mu, logvar, discriminator, lpips_fn, beta_kld=1.0)
                     # loss, *_ = vae_loss(recon_x, x, mu, logvar, lpips_fn, ssim_fn, beta_kld=vae_beta_kld)
-                    loss, *_ = vae_loss(recon_x, x, mu, logvar, beta_kld=vae_beta_kld)
                 test_loss += loss.item()
             avg_test_loss = test_loss / len(test_loader)
             scheduler.step(avg_test_loss)
@@ -187,7 +186,6 @@ def plot_recon_vs_kld(train_bce_loss, train_kld_loss):
     plt.figure(figsize=(10, 6))
     plt.plot(epochs, train_bce_loss, label='Reconstruction Loss (BCE)', color='blue', linewidth=2)
     plt.plot(epochs, train_kld_loss, label='KL Divergence', color='red', linewidth=2)
-    # plt.plot(epochs, train_tvl_loss, label='TVL', color='yellow', linewidth=2)
     plt.xlabel('Epoch')
     plt.ylabel('Loss')
     plt.title('Reconstruction Loss vs KL Divergence over Epochs')

 import sys, os
 sys.path.insert(0, os.path.dirname(__file__))
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 import torch
 from torch import nn, optim
 import torch.nn.functional as F
 import lpips
 from pytorch_msssim import SSIM
 # # Simple PatchGAN Discriminator
 # class PatchDiscriminator(nn.Module):
     tvl_w = torch.pow(x[:, :, :, 1:] - x[:, :, :, :-1], 2).sum() # TV-L2
     return (tvl_h + tvl_w) / x.size(0)
+def get_annealed_beta(epoch, warmup_epochs=100, max_beta=1.0): return vae_beta_kld * min(max_beta, epoch / warmup_epochs)
 # def vae_loss(recon_x, x, mu, logvar, lpips_fn, ssim_fn, beta_kld=1.0):
+def vae_loss(recon_x, x, mu, logvar, beta_kld=1.0, lpips_fn=None, ssim_fn=None):
     # mse = torch.mean((recon_x - x) ** 2)
     # kld_loss = torch.mean(-0.5 * (1 + logvar - mu.pow(2) - logvar.exp()))
     # b, c, h, w = x.shape
+    tvl = total_variance_loss(recon_x) if vae_lambda_tvl > 0 else 0.0
     mse = F.mse_loss(recon_x, x, reduction="sum") / recon_x.size(0)
     kld_loss = torch.sum(-0.5 * (1 + logvar - mu.pow(2) - logvar.exp())) / recon_x.size(0)
+    # mse = F.mse_loss(recon_x, x, reduction="mean")
+    # kld_loss = torch.mean(-0.5 * (1 + logvar - mu.pow(2) - logvar.exp()))  # divides by latent_ch×H×W×B
+    with torch.amp.autocast("cuda", enabled=False):
+        ssim_loss = 1 - ssim_fn(recon_x.float(), x.float())
+        lpips_loss = lpips_fn(recon_x.float() * 2 - 1, x.float() * 2 - 1).mean()
+    total_loss = mse + (beta_kld * kld_loss) + (vae_lambda_tvl * tvl) + (vae_lpips_weight * lpips_loss) + ssim_loss * 0.1
+    return total_loss, mse, kld_loss, tvl, lpips_loss, ssim_loss
+    # total_loss = mse + (beta_kld * kld_loss)
+    # return total_loss, mse, kld_loss, 0, 0, 0
 def train_test_vae():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5)
     os.makedirs(f"{vae_checkpoint_dir}", exist_ok=True)
     early_stopping_counter = 0
     best_test_loss = float("inf")
     train_bce_loss = []
     train_kld_loss = []
+    lpips_fn = lpips.LPIPS(net='vgg').to(device)
+    lpips_fn.eval()
+    for param in lpips_fn.parameters():
+        param.requires_grad = False
+    ssim_fn = SSIM(data_range=1, size_average=True, channel=3)
     scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
     # discriminator = PatchDiscriminator().to(device)
     # optimizer_D = optim.AdamW(discriminator.parameters(), lr=vae_optim_lr, betas=(0.5, 0.999))
                 # loss, rec, kld, tvl, lpips_loss, gan_loss = vae_loss(recon_x, x, mu, logvar, discriminator, lpips_fn, beta_kld=get_annealed_beta(epoch)) # Update VAE (generator)
                 # loss, rec, kld, tvl, lpips_loss = vae_loss(recon_x, x, mu, logvar, lpips_fn, beta_kld=min(1, global_step / warmup_steps)) # Update VAE (generator)
                 # loss, rec, kld, tvl, lpips_loss, ssim_loss = vae_loss(recon_x, x, mu, logvar, lpips_fn, ssim_fn, beta_kld=get_annealed_beta(epoch)) # Update VAE (generator)
+                loss, rec, kld, tvl, lpips_loss, ssim_loss = vae_loss(recon_x, x, mu, logvar, beta_kld=get_annealed_beta(epoch), lpips_fn=lpips_fn, ssim_fn=ssim_fn) # Update VAE (generator)
             optimizer.zero_grad(set_to_none=True)
             scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(vae.parameters(), max_norm=1.0)
             scaler.step(optimizer)
             scaler.update()
             train_loss += loss.item()
             total_rec += rec.item()
             total_kld += kld.item()
+            total_tvl += tvl.item() if vae_lambda_tvl > 0 else 0.0
+            total_lpips += lpips_loss.item()
+            total_ssim += ssim_loss.item()
             # total_gan += gan_loss.item()
         avg_train_loss = train_loss / len(train_loader)
         avg_rec = total_rec / len(train_loader)
                     recon_x, mu, logvar = vae(x)
                     # loss, *_ = vae_loss(recon_x, x, mu, logvar, discriminator, lpips_fn, beta_kld=1.0)
                     # loss, *_ = vae_loss(recon_x, x, mu, logvar, lpips_fn, ssim_fn, beta_kld=vae_beta_kld)
+                    loss, *_ = vae_loss(recon_x, x, mu, logvar, beta_kld=vae_beta_kld, lpips_fn=lpips_fn, ssim_fn=ssim_fn)
                 test_loss += loss.item()
             avg_test_loss = test_loss / len(test_loader)
             scheduler.step(avg_test_loss)
     plt.figure(figsize=(10, 6))
     plt.plot(epochs, train_bce_loss, label='Reconstruction Loss (BCE)', color='blue', linewidth=2)
     plt.plot(epochs, train_kld_loss, label='KL Divergence', color='red', linewidth=2)
     plt.xlabel('Epoch')
     plt.ylabel('Loss')
     plt.title('Reconstruction Loss vs KL Divergence over Epochs')

core/unet.py CHANGED Viewed

@@ -12,8 +12,8 @@ def get_timestep_embedding(timesteps: torch.Tensor, dim: int = 1024):
     """Sinusoidal position embeddings"""
     assert timesteps.dim() == 1
     half_dim = dim // 2
-    freqs = torch.exp(-math.log(10000) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / (half_dim - 1))
-    # freqs = torch.exp(-math.log(10000) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / (half_dim))
     args = timesteps.float().unsqueeze(1) * freqs.unsqueeze(0)
     emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
     if dim % 2: emb = F.pad(emb, (0, 1))  # zero‑pad if dim is odd
@@ -32,7 +32,7 @@ class ResBlock(nn.Module):
         super().__init__()
         self.convgnact1 = ConvGNAct(in_ch, out_ch)
         self.norm = nn.GroupNorm(unet_group_size, out_ch)
-        self.act = nn.SiLU(inplace=True)
         self.out_conv = nn.Conv2d(out_ch, out_ch, 3, padding=1)
         self.time_emb_proj = nn.Sequential(
             nn.SiLU(),
@@ -60,26 +60,17 @@ class DownBlock(nn.Module):
         super().__init__()
         self.resblock1 = ResBlock(in_ch, out_ch, time_emb_dim)
         self.resblock2 = ResBlock(out_ch, out_ch, time_emb_dim)
-        self.resblock3 = ResBlock(out_ch, out_ch, time_emb_dim)
-        self.resblock4 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock5 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock6 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock7 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock8 = ResBlock(out_ch, out_ch, time_emb_dim)
-        self.attn = SelfCrossAttn(out_ch, heads=heads, text_emb_dim=text_emb_dim, cross=cross)
         # self.down = ConvGNAct(out_ch, out_ch, 3, 2, 1)
         self.down = nn.Conv2d(out_ch, out_ch, 3, 2, 1)
         self.use_attn = use_attn
     def forward(self, x, time_emb, text):
         x = self.resblock1(x, time_emb)
         x = self.resblock2(x, time_emb)
-        x = self.resblock3(x, time_emb)
-        x = self.resblock4(x, time_emb)
-        # x = self.resblock5(x, time_emb)
-        # x = self.resblock6(x, time_emb)
-        # x = self.resblock7(x, time_emb)
-        # x = self.resblock8(x, time_emb)
-        if self.use_attn: x = self.attn(x, text)
         x_down = self.down(x)
         return x_down, x
 class UpBlock(nn.Module):
@@ -88,13 +79,8 @@ class UpBlock(nn.Module):
         # self.up = nn.Conv2d(in_ch, in_ch, kernel_size=3, padding=1)
         self.resblock1 = ResBlock(in_ch * 2, out_ch, time_emb_dim)
         self.resblock2 = ResBlock(out_ch, out_ch, time_emb_dim)
-        self.resblock3 = ResBlock(out_ch, out_ch, time_emb_dim)
-        self.resblock4 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock5 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock6 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock7 = ResBlock(out_ch, out_ch, time_emb_dim)
-        # self.resblock8 = ResBlock(out_ch, out_ch, time_emb_dim)
-        self.attn = SelfCrossAttn(out_ch, heads=heads, text_emb_dim=text_emb_dim, cross=cross)
         # self.out_proj = ConvGNAct(out_ch, out_ch, k=1, s=1, p=0)
         self.use_attn = use_attn
     def forward(self, x, skip, time_emb, text):
@@ -104,13 +90,9 @@ class UpBlock(nn.Module):
         x = torch.cat([x, skip], dim=1)
         x = self.resblock1(x, time_emb)
         x = self.resblock2(x, time_emb)
-        x = self.resblock3(x, time_emb)
-        x = self.resblock4(x, time_emb)
-        # x = self.resblock5(x, time_emb)
-        # x = self.resblock6(x, time_emb)
-        # x = self.resblock7(x, time_emb)
-        # x = self.resblock8(x, time_emb)
-        if self.use_attn: x = self.attn(x, text)
         # x = self.out_proj(x)
         return x
 class Unet(nn.Module):
@@ -125,13 +107,14 @@ class Unet(nn.Module):
         # Downsample
         self.down1 = DownBlock(in_ch=base_ch, out_ch=base_ch * 2, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=False, heads=4)
         self.down2 = DownBlock(in_ch=base_ch * 2, out_ch=base_ch * 4, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=8)
-        self.down3 = DownBlock(in_ch=base_ch * 4, out_ch=base_ch * 8, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=8)
         # Mid (bottleneck)
         self.mid_res1 = ResBlock(in_ch=base_ch * 8, out_ch=base_ch * 8, time_emb_dim=time_emb_dim)
-        self.mid_attn = SelfCrossAttn(base_ch * 8, heads=16, text_emb_dim=1024, cross=True)
         self.mid_res2 = ResBlock(in_ch=base_ch * 8, out_ch=base_ch * 8, time_emb_dim=time_emb_dim)
         # Upsample
-        self.up3 = UpBlock(in_ch=base_ch * 8, out_ch=base_ch * 4, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=8)
         self.up2 = UpBlock(in_ch=base_ch * 4, out_ch=base_ch * 2, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=8)
         self.up1 = UpBlock(in_ch=base_ch * 2, out_ch=base_ch, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=False, heads=4)
         self.out_norm = nn.GroupNorm(unet_group_size, base_ch)
@@ -144,7 +127,8 @@ class Unet(nn.Module):
         h, skip2 = self.down2(h, time_emb, text_emb) # (B, C x 4, H / 4, W / 4) (B, C x 4, H / 2, W / 2) # print("DOWN 2", h.shape, skip2.shape)
         h, skip3 = self.down3(h, time_emb, text_emb) # (B, C x 8, H / 8, W / 8) (B, C x 4, H / 4, W / 4) # print("DOWN 3", h.shape, skip3.shape)
         h = self.mid_res1(h, time_emb) # (B, C x 8, H / 8, W / 8) # print("MID RES", h.shape)
-        h = self.mid_attn(h, text_emb) # (B, C x 8, H / 8, W / 8) # print("MID ATTN", h.shape)
         h = self.mid_res2(h, time_emb) # (B, C x 8, H / 8, W / 8) # print("MID RES", h.shape)
         h = self.up3(h, skip3, time_emb, text_emb) # (B, C x 4, H / 4, W / 4) # print("UP 2", h.shape)
         h = self.up2(h, skip2, time_emb, text_emb) # (B, C x 2, H / 2, W / 2) # print("UP 2", h.shape)

     """Sinusoidal position embeddings"""
     assert timesteps.dim() == 1
     half_dim = dim // 2
+    # freqs = torch.exp(-math.log(10000) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / (half_dim - 1))
+    freqs = torch.exp(-math.log(10000) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / (half_dim))
     args = timesteps.float().unsqueeze(1) * freqs.unsqueeze(0)
     emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
     if dim % 2: emb = F.pad(emb, (0, 1))  # zero‑pad if dim is odd
         super().__init__()
         self.convgnact1 = ConvGNAct(in_ch, out_ch)
         self.norm = nn.GroupNorm(unet_group_size, out_ch)
+        self.act = nn.SiLU()
         self.out_conv = nn.Conv2d(out_ch, out_ch, 3, padding=1)
         self.time_emb_proj = nn.Sequential(
             nn.SiLU(),
         super().__init__()
         self.resblock1 = ResBlock(in_ch, out_ch, time_emb_dim)
         self.resblock2 = ResBlock(out_ch, out_ch, time_emb_dim)
+        self.self_attn = SelfCrossAttn(out_ch, heads=heads, text_emb_dim=text_emb_dim, cross=False)
+        self.cross_attn = SelfCrossAttn(out_ch, heads=heads, text_emb_dim=text_emb_dim, cross=True)
         # self.down = ConvGNAct(out_ch, out_ch, 3, 2, 1)
         self.down = nn.Conv2d(out_ch, out_ch, 3, 2, 1)
         self.use_attn = use_attn
     def forward(self, x, time_emb, text):
         x = self.resblock1(x, time_emb)
         x = self.resblock2(x, time_emb)
+        if self.use_attn:
+            x = self.self_attn(x, text)
+            x = self.cross_attn(x, text)
         x_down = self.down(x)
         return x_down, x
 class UpBlock(nn.Module):
         # self.up = nn.Conv2d(in_ch, in_ch, kernel_size=3, padding=1)
         self.resblock1 = ResBlock(in_ch * 2, out_ch, time_emb_dim)
         self.resblock2 = ResBlock(out_ch, out_ch, time_emb_dim)
+        self.self_attn = SelfCrossAttn(out_ch, heads=heads, text_emb_dim=text_emb_dim, cross=False)
+        self.cross_attn = SelfCrossAttn(out_ch, heads=heads, text_emb_dim=text_emb_dim, cross=True)
         # self.out_proj = ConvGNAct(out_ch, out_ch, k=1, s=1, p=0)
         self.use_attn = use_attn
     def forward(self, x, skip, time_emb, text):
         x = torch.cat([x, skip], dim=1)
         x = self.resblock1(x, time_emb)
         x = self.resblock2(x, time_emb)
+        if self.use_attn:
+            x = self.self_attn(x, text)
+            x = self.cross_attn(x, text)
         # x = self.out_proj(x)
         return x
 class Unet(nn.Module):
         # Downsample
         self.down1 = DownBlock(in_ch=base_ch, out_ch=base_ch * 2, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=False, heads=4)
         self.down2 = DownBlock(in_ch=base_ch * 2, out_ch=base_ch * 4, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=8)
+        self.down3 = DownBlock(in_ch=base_ch * 4, out_ch=base_ch * 8, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=16)
         # Mid (bottleneck)
         self.mid_res1 = ResBlock(in_ch=base_ch * 8, out_ch=base_ch * 8, time_emb_dim=time_emb_dim)
+        self.mid_self_attn = SelfCrossAttn(base_ch * 8, heads=32, text_emb_dim=1024, cross=False)
+        self.mid_cross_attn = SelfCrossAttn(base_ch * 8, heads=32, text_emb_dim=1024, cross=True)
         self.mid_res2 = ResBlock(in_ch=base_ch * 8, out_ch=base_ch * 8, time_emb_dim=time_emb_dim)
         # Upsample
+        self.up3 = UpBlock(in_ch=base_ch * 8, out_ch=base_ch * 4, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=16)
         self.up2 = UpBlock(in_ch=base_ch * 4, out_ch=base_ch * 2, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=True, heads=8)
         self.up1 = UpBlock(in_ch=base_ch * 2, out_ch=base_ch, time_emb_dim=time_emb_dim, text_emb_dim=text_emb_dim, cross=True, use_attn=False, heads=4)
         self.out_norm = nn.GroupNorm(unet_group_size, base_ch)
         h, skip2 = self.down2(h, time_emb, text_emb) # (B, C x 4, H / 4, W / 4) (B, C x 4, H / 2, W / 2) # print("DOWN 2", h.shape, skip2.shape)
         h, skip3 = self.down3(h, time_emb, text_emb) # (B, C x 8, H / 8, W / 8) (B, C x 4, H / 4, W / 4) # print("DOWN 3", h.shape, skip3.shape)
         h = self.mid_res1(h, time_emb) # (B, C x 8, H / 8, W / 8) # print("MID RES", h.shape)
+        h = self.mid_self_attn(h, text_emb) # (B, C x 8, H / 8, W / 8) # print("MID ATTN", h.shape)
+        h = self.mid_cross_attn(h, text_emb) # (B, C x 8, H / 8, W / 8) # print("MID ATTN", h.shape)
         h = self.mid_res2(h, time_emb) # (B, C x 8, H / 8, W / 8) # print("MID RES", h.shape)
         h = self.up3(h, skip3, time_emb, text_emb) # (B, C x 4, H / 4, W / 4) # print("UP 2", h.shape)
         h = self.up2(h, skip2, time_emb, text_emb) # (B, C x 2, H / 2, W / 2) # print("UP 2", h.shape)

core/vae.py CHANGED Viewed

@@ -6,32 +6,50 @@ import torch.nn.functional as F
 from config import *
 from attention import SelfCrossAttn
 class VAE(nn.Module):
     def __init__(self):
         super().__init__()
         # Encoder
         self.encoder_conv = nn.Sequential(
-            nn.Conv2d(3, 32, 3, 1, padding=1, bias=False), # (B, 4, 128, 128)
             nn.GroupNorm(vae_group_size, 32), nn.SiLU(inplace=True),
-            nn.Conv2d(32, 128, 3, 2, padding=1, bias=False), # (B, 32, 64, 64)
             nn.GroupNorm(vae_group_size, 128), nn.SiLU(inplace=True),
-            nn.Conv2d(128, 256, 3, 2, padding=1, bias=False), # (B, 128, 32, 32)
             nn.GroupNorm(vae_group_size, 256), nn.SiLU(inplace=True),
             # nn.Conv2d(256, 512, 3, 2, padding=1, bias=False), # (B, 128, 16, 16)
             # nn.GroupNorm(vae_group_size, 512), nn.SiLU(inplace=True),
-            SelfCrossAttn(256, heads=8, cross=False),
         )
-        # 1×1×1 convs give channel‑wise μ and log σ², shape = (B, latent_channels, 4, 4, 4)
-        self.to_latent = nn.Conv2d(256, 2 * vae_latent_channels, kernel_size=1)
-        self.conv_mu = nn.Conv2d(vae_latent_channels, vae_latent_channels, kernel_size=1)
-        self.conv_logvar = nn.Conv2d(vae_latent_channels, vae_latent_channels, kernel_size=1)
-        self.from_latent = nn.Conv2d(vae_latent_channels, 256, kernel_size=1)
         # Decoder
         self.decoder_conv = nn.Sequential(
-            SelfCrossAttn(256, heads=8, cross=False),
-            # nn.Upsample(scale_factor=2, mode="nearest"),
-            # nn.Conv2d(512, 256, 3, padding=1, bias=False),
-            # nn.GroupNorm(vae_group_size, 256), nn.SiLU(inplace=True),
             nn.Upsample(scale_factor=2, mode="nearest"),
             nn.Conv2d(256, 128, 3, padding=1, bias=False),
             nn.GroupNorm(vae_group_size, 128), nn.SiLU(inplace=True),
@@ -47,13 +65,12 @@ class VAE(nn.Module):
         std = torch.exp(0.5 * logvar)
         eps = torch.randn_like(std)
         return mu + eps * std
     def forward(self, x):
         h = self.encoder_conv(x) # (B, C, D, H)
         h = self.to_latent(h)
-        # mu, logvar = self.conv_mu(h), self.conv_logvar(h)
         mu, logvar = torch.chunk(h, 2, dim=1)
         z = self.reparameterize(mu, logvar) # Latent (B, C, D, H)
-        # z = z * 0.18215
         h = self.from_latent(z)
         recon = self.decoder_conv(h)
         return recon, mu, logvar
@@ -61,7 +78,6 @@ class VAE(nn.Module):
     def encode_img_to_latent(self, x):
         h = self.encoder_conv(x) # (B, C, D, H)
         h = self.to_latent(h)
-        # mu, logvar = self.conv_mu(h), self.conv_logvar(h)
         mu, logvar = torch.chunk(h, 2, dim=1)
         z = self.reparameterize(mu, logvar)
         return z

 from config import *
 from attention import SelfCrossAttn
+class VAEResBlock(nn.Module):
+    def __init__(self, in_channels, out_channels=None):
+        out_channels = out_channels or in_channels
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.GroupNorm(vae_group_size, in_channels), nn.SiLU(inplace=True),
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.GroupNorm(vae_group_size, out_channels), nn.SiLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+        )
+        self.skip = nn.Conv2d(in_channels, out_channels, 1, bias=False) if in_channels != out_channels else nn.Identity()
+    def forward(self, x):
+        return self.block(x) + self.skip(x)
 class VAE(nn.Module):
     def __init__(self):
         super().__init__()
         # Encoder
         self.encoder_conv = nn.Sequential(
+            nn.Conv2d(3, 32, 3, 1, padding=1, bias=False), # (B, 4, 256, 256)
             nn.GroupNorm(vae_group_size, 32), nn.SiLU(inplace=True),
+            nn.Conv2d(32, 128, 3, 2, padding=1, bias=False), # (B, 32, 128, 128)
             nn.GroupNorm(vae_group_size, 128), nn.SiLU(inplace=True),
+            nn.Conv2d(128, 256, 3, 2, padding=1, bias=False), # (B, 128, 64, 64)
             nn.GroupNorm(vae_group_size, 256), nn.SiLU(inplace=True),
+            nn.Conv2d(256, 512, 3, 2, padding=1, bias=False), # (B, 128, 32, 32)
+            nn.GroupNorm(vae_group_size, 512), nn.SiLU(inplace=True),
             # nn.Conv2d(256, 512, 3, 2, padding=1, bias=False), # (B, 128, 16, 16)
             # nn.GroupNorm(vae_group_size, 512), nn.SiLU(inplace=True),
+            # SelfCrossAttn(512, heads=8, cross=False),
+            VAEResBlock(512), SelfCrossAttn(512, heads=8, cross=False), VAEResBlock(512),
+            nn.GroupNorm(vae_group_size, 512), nn.SiLU(inplace=True),
         )
+        # Channel‑wise μ and log σ², shape = (B, latent_channels, 4, 4, 4)
+        self.to_latent = nn.Conv2d(512, 2 * vae_latent_channels, kernel_size=1)
+        self.from_latent = nn.Conv2d(vae_latent_channels, 512, kernel_size=1)
         # Decoder
         self.decoder_conv = nn.Sequential(
+            VAEResBlock(512), SelfCrossAttn(512, heads=8, cross=False), VAEResBlock(512),
+            nn.GroupNorm(vae_group_size, 512), nn.SiLU(inplace=True),
+            # SelfCrossAttn(512, heads=8, cross=False),
+            nn.Upsample(scale_factor=2, mode="nearest"),
+            nn.Conv2d(512, 256, 3, padding=1, bias=False),
+            nn.GroupNorm(vae_group_size, 256), nn.SiLU(inplace=True),
             nn.Upsample(scale_factor=2, mode="nearest"),
             nn.Conv2d(256, 128, 3, padding=1, bias=False),
             nn.GroupNorm(vae_group_size, 128), nn.SiLU(inplace=True),
         std = torch.exp(0.5 * logvar)
         eps = torch.randn_like(std)
         return mu + eps * std
     def forward(self, x):
         h = self.encoder_conv(x) # (B, C, D, H)
         h = self.to_latent(h)
         mu, logvar = torch.chunk(h, 2, dim=1)
         z = self.reparameterize(mu, logvar) # Latent (B, C, D, H)
         h = self.from_latent(z)
         recon = self.decoder_conv(h)
         return recon, mu, logvar
     def encode_img_to_latent(self, x):
         h = self.encoder_conv(x) # (B, C, D, H)
         h = self.to_latent(h)
         mu, logvar = torch.chunk(h, 2, dim=1)
         z = self.reparameterize(mu, logvar)
         return z