Create critic.py

jobs/process/models/critic.py

@@ -0,0 +1,234 @@
+import glob
+import os
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file, save_file
+
+from toolkit.losses import get_gradient_penalty
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.optimizer import get_optimizer
+from toolkit.train_tools import get_torch_dtype
+
+
+class MeanReduce(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, inputs):
+        # global mean over spatial dims (keeps channel/batch)
+        return torch.mean(inputs, dim=(2, 3), keepdim=True)
+
+
+class SelfAttention2d(nn.Module):
+    """
+    Lightweight self-attention layer (SAGAN-style) that keeps spatial
+    resolution unchanged. Adds minimal params / compute but improves
+    long-range modelling – helpful for variable-sized inputs.
+    """
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.query = nn.Conv1d(in_channels, in_channels // 8, 1)
+        self.key   = nn.Conv1d(in_channels, in_channels // 8, 1)
+        self.value = nn.Conv1d(in_channels, in_channels,      1)
+        self.gamma = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        flat = x.view(B, C, H * W)                    # (B,C,N)
+        q = self.query(flat).permute(0, 2, 1)         # (B,N,C//8)
+        k = self.key(flat)                            # (B,C//8,N)
+        attn = torch.bmm(q, k)                        # (B,N,N)
+        attn = attn.softmax(dim=-1)                   # softmax along last dim
+        v = self.value(flat)                          # (B,C,N)
+        out = torch.bmm(v, attn.permute(0, 2, 1))     # (B,C,N)
+        out = out.view(B, C, H, W)                    # restore spatial dims
+        return self.gamma * out + x                   # residual
+
+
+class CriticModel(nn.Module):
+    def __init__(self, base_channels: int = 64):
+        super().__init__()
+
+        def sn_conv(in_c, out_c, k, s, p):
+            return nn.utils.spectral_norm(
+                nn.Conv2d(in_c, out_c, kernel_size=k, stride=s, padding=p)
+            )
+
+        layers = [
+            # initial down-sample
+            sn_conv(3, base_channels, 3, 2, 1),
+            nn.LeakyReLU(0.2, inplace=True),
+        ]
+
+        in_c = base_channels
+        # progressive downsamples ×3 (64→128→256→512)
+        for _ in range(3):
+            out_c = min(in_c * 2, 1024)
+            layers += [
+                sn_conv(in_c, out_c, 3, 2, 1),
+                nn.LeakyReLU(0.2, inplace=True),
+            ]
+            # single attention block after reaching 256 channels
+            if out_c == 256:
+                layers += [SelfAttention2d(out_c)]
+            in_c = out_c
+
+        # extra depth (keeps spatial size)
+        layers += [
+            sn_conv(in_c, 1024, 3, 1, 1),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            # final 1-channel prediction map
+            sn_conv(1024, 1, 3, 1, 1),
+            MeanReduce(),        # → (B,1,1,1)
+            nn.Flatten(),        # → (B,1)
+        ]
+
+        self.main = nn.Sequential(*layers)
+
+    def forward(self, inputs):
+        # force full-precision inside AMP ctx for stability
+        with torch.cuda.amp.autocast(False):
+            return self.main(inputs.float())
+
+
+if TYPE_CHECKING:
+    from jobs.process.TrainVAEProcess import TrainVAEProcess
+    from jobs.process.TrainESRGANProcess import TrainESRGANProcess
+
+
+class Critic:
+    process: Union['TrainVAEProcess', 'TrainESRGANProcess']
+
+    def __init__(
+        self,
+        learning_rate=1e-5,
+        device='cpu',
+        optimizer='adam',
+        num_critic_per_gen=1,
+        dtype='float32',
+        lambda_gp=10,
+        start_step=0,
+        warmup_steps=1000,
+        process=None,
+        optimizer_params=None,
+    ):
+        self.learning_rate = learning_rate
+        self.device = device
+        self.optimizer_type = optimizer
+        self.num_critic_per_gen = num_critic_per_gen
+        self.dtype = dtype
+        self.torch_dtype = get_torch_dtype(self.dtype)
+        self.process = process
+        self.model = None
+        self.optimizer = None
+        self.scheduler = None
+        self.warmup_steps = warmup_steps
+        self.start_step = start_step
+        self.lambda_gp = lambda_gp
+
+        if optimizer_params is None:
+            optimizer_params = {}
+        self.optimizer_params = optimizer_params
+        self.print = self.process.print
+        print(f" Critic config: {self.__dict__}")
+
+    def setup(self):
+        self.model = CriticModel().to(self.device)
+        self.load_weights()
+        self.model.train()
+        self.model.requires_grad_(True)
+        params = self.model.parameters()
+        self.optimizer = get_optimizer(
+            params,
+            self.optimizer_type,
+            self.learning_rate,
+            optimizer_params=self.optimizer_params,
+        )
+        self.scheduler = torch.optim.lr_scheduler.ConstantLR(
+            self.optimizer,
+            total_iters=self.process.max_steps * self.num_critic_per_gen,
+            factor=1,
+            verbose=False,
+        )
+
+    def load_weights(self):
+        path_to_load = None
+        self.print(f"Critic: Looking for latest checkpoint in {self.process.save_root}")
+        files = glob.glob(os.path.join(self.process.save_root, f"CRITIC_{self.process.job.name}*.safetensors"))
+        if files:
+            latest_file = max(files, key=os.path.getmtime)
+            print(f" - Latest checkpoint is: {latest_file}")
+            path_to_load = latest_file
+        else:
+            self.print(" - No checkpoint found, starting from scratch")
+        if path_to_load:
+            self.model.load_state_dict(load_file(path_to_load))
+
+    def save(self, step=None):
+        self.process.update_training_metadata()
+        save_meta = get_meta_for_safetensors(self.process.meta, self.process.job.name)
+        step_num = f"_{str(step).zfill(9)}" if step is not None else ''
+        save_path = os.path.join(
+            self.process.save_root, f"CRITIC_{self.process.job.name}{step_num}.safetensors"
+        )
+        save_file(self.model.state_dict(), save_path, save_meta)
+        self.print(f"Saved critic to {save_path}")
+
+    def get_critic_loss(self, vgg_output):
+        # (caller still passes combined [pred|target] images)
+        if self.start_step > self.process.step_num:
+            return torch.tensor(0.0, dtype=self.torch_dtype, device=self.device)
+
+        warmup_scaler = 1.0
+        if self.process.step_num < self.start_step + self.warmup_steps:
+            warmup_scaler = (self.process.step_num - self.start_step) / self.warmup_steps
+
+        self.model.eval()
+        self.model.requires_grad_(False)
+
+        vgg_pred, _ = torch.chunk(vgg_output.float(), 2, dim=0)
+        stacked_output = self.model(vgg_pred)
+        return (-torch.mean(stacked_output)) * warmup_scaler
+
+    def step(self, vgg_output):
+        self.model.train()
+        self.model.requires_grad_(True)
+        self.optimizer.zero_grad()
+
+        critic_losses = []
+        inputs = vgg_output.detach().to(self.device, dtype=torch.float32)
+
+        vgg_pred, vgg_target = torch.chunk(inputs, 2, dim=0)
+        stacked_output = self.model(inputs).float()
+        out_pred, out_target = torch.chunk(stacked_output, 2, dim=0)
+
+        # hinge loss + gradient penalty
+        loss_real = torch.relu(1.0 - out_target).mean()
+        loss_fake = torch.relu(1.0 + out_pred).mean()
+        gradient_penalty = get_gradient_penalty(self.model, vgg_target, vgg_pred, self.device)
+        critic_loss = loss_real + loss_fake + self.lambda_gp * gradient_penalty
+
+        critic_loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+        self.optimizer.step()
+        self.scheduler.step()
+        critic_losses.append(critic_loss.item())
+
+        return float(np.mean(critic_losses))
+
+    def get_lr(self):    
+        if hasattr(self.optimizer, 'get_avg_learning_rate'):
+            learning_rate = self.optimizer.get_avg_learning_rate()
+        elif self.optimizer_type.startswith('dadaptation') or \
+                self.optimizer_type.lower().startswith('prodigy'):
+            learning_rate = (
+                self.optimizer.param_groups[0]["d"] *
+                self.optimizer.param_groups[0]["lr"]
+            )
+        else:
+            learning_rate = self.optimizer.param_groups[0]['lr']
+        return learning_rate