Spaces:
Running
Running
| 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 |