File size: 8,447 Bytes
d446991
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
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