1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "start_line": 32,
5  "end_line": 76,
6  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Decomposed Multi-Band VAE Loss.\n\n    Decomposes images into low/mid/high frequency bands using radial\n    butterworth filters in the Fourier domain. Each band receives a\n    specialized loss: L1 for structure (low), LPIPS for perception (mid),\n    and gradient matching for fine details (high). Progressive scheduling\n    shifts emphasis from coarse to fine as training proceeds.\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.total_steps = 5000\n\n        # Precompute frequency band masks for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n\n        # Butterworth filter: 1 / (1 + (r / r_cutoff)^(2n))\n        n_order = 2\n        low_r = 0.25\n        high_r = 0.55\n        low_pass = 1.0 / (1.0 + (r_norm / low_r) ** (2 * n_order))\n        high_pass = 1.0 / (1.0 + (high_r / r_norm) ** (2 * n_order))\n        mid_pass = (1.0 - low_pass) * high_pass\n\n        # Reshape for broadcasting with [B, C, H, W]\n        self.register_buffer('low_mask', low_pass.view(1, 1, H, W))\n        self.register_buffer('mid_mask', mid_pass.view(1, 1, H, W))\n        self.register_buffer('high_mask', (1.0 - high_pass).view(1, 1, H, W))\n\n    def _fft_filter(self, img, mask):\n        \"\"\"Apply frequency-domain filter.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude using finite differences with padding.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # Decompose into frequency bands\n        recon_low = self._fft_filter(recon, self.low_mask)\n        recon_mid = self._fft_filter(recon, self.mid_mask)\n        recon_high = self._fft_filter(recon, self.high_mask)\n\n        target_low = self._fft_filter(target, self.low_mask)\n        target_mid = self._fft_filter(target, self.mid_mask)\n        target_high = self._fft_filter(target, self.high_mask)\n\n        # L1 loss on low-frequency (structure)\n        rec_loss_low = F.l1_loss(recon_low, target_low)\n\n        # LPIPS loss on mid-frequency (perceptual features)\n        p_loss = self.lpips_fn(recon_mid.float(), target_mid.float()).mean()\n\n        # Gradient matching on high-frequency (edge sharpness)\n        recon_grad = self._grad_magnitude(recon_high)\n        target_grad = self._grad_magnitude(target_high)\n        grad_loss = F.l1_loss(recon_grad, target_grad)\n\n        # Progressive frequency scheduling: coarse-to-fine training\n        progress = min(step / self.total_steps, 1.0)\n        # Low freq: strong early, moderate late\n        w_low = 1.0 - 0.4 * progress\n        # Mid freq: moderate throughout, slight increase\n        w_mid = 0.5 + 0.3 * progress\n        # High freq: weak early, strong late (fine details)\n        w_high = 0.1 + 0.5 * progress\n\n        # KL divergence\n        kl_loss = posterior.kl().mean()\n\n        # Total loss\n        loss = (w_low * rec_loss_low +\n                w_mid * p_loss +\n                w_high * grad_loss +\n                self.kl_weight * kl_loss)\n\n        return loss, {\n            \"rec_loss_low\": rec_loss_low.item(),\n            \"p_loss\": p_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"w_low\": w_low,\n            \"w_mid\": w_mid,\n            \"w_high\": w_high,\n        }"
7}

1{
2  "op": "replace",
3  "filename": "diffusers-main/custom_train.py",
4  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Decomposed Adversarial VAE Loss.\n\n    Combines frequency-band decomposition with adversarial training:\n    - Low-band: L1 reconstruction (structure preservation)\n    - Mid-band: LPIPS perceptual loss (feature matching)\n    - High-band: gradient matching + adversarial for sharpness\n    - Discriminator with feature matching loss\n    - Adaptive weighting from taming-transformers style\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.perceptual_weight = 0.5\n\n        # Precompute frequency band masks for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n\n        # Butterworth filter order 2\n        n_order = 2\n        low_r = 0.25\n        high_r = 0.55\n        low_pass = 1.0 / (1.0 + (r_norm / low_r) ** (2 * n_order))\n        high_pass = 1.0 / (1.0 + (high_r / r_norm) ** (2 * n_order))\n        mid_pass = (1.0 - low_pass) * high_pass\n\n        self.register_buffer('low_mask', low_pass.view(1, 1, H, W))\n        self.register_buffer('mid_mask', mid_pass.view(1, 1, H, W))\n        self.register_buffer('high_mask', (1.0 - high_pass).view(1, 1, H, W))\n\n        # Build discriminator with spectral norm and feature extraction\n        from torch.nn.utils import spectral_norm\n        self.disc = _build_discriminator()\n        for m in self.disc.modules():\n            if isinstance(m, nn.Conv2d):\n                self.disc._modules[m] = spectral_norm(m)\n        # Apply spectral norm properly\n        self.disc = _build_disc_sn()\n        self.disc.to(device)\n        self.disc_opt = torch.optim.Adam(self.disc.parameters(), lr=1e-4, betas=(0.5, 0.9))\n        self.disc_start = 3000\n        self.feat_match_weight = 1.0\n        self.gan_weight = 0.8\n        self.grad_weight = 0.3\n\n    def _fft_filter(self, img, mask):\n        \"\"\"Apply frequency-domain filter.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def _compute_gan_and_feat_loss(self, recon, target):\n        \"\"\"Compute adversarial loss and feature matching loss.\"\"\"\n        # Generator loss: fool discriminator\n        fake_pred, fake_feats = self.disc(recon, return_features=True)\n        gen_loss = -fake_pred.mean()\n\n        # Feature matching loss\n        with torch.no_grad():\n            _, real_feats = self.disc(target, return_features=True)\n        feat_loss = 0.0\n        for rf, ff in zip(real_feats, fake_feats):\n            feat_loss += F.l1_loss(ff, rf)\n\n        # Discriminator loss\n        real_pred, _ = self.disc(target, return_features=True)\n        fake_pred_d, _ = self.disc(recon.detach(), return_features=True)\n        disc_loss = F.relu(1.0 - real_pred).mean() + F.relu(1.0 + fake_pred_d).mean()\n\n        return gen_loss, feat_loss, disc_loss\n\n    def forward(self, recon, target, posterior, step):\n        # Frequency decomposition\n        recon_low = self._fft_filter(recon, self.low_mask)\n        recon_mid = self._fft_filter(recon, self.mid_mask)\n        recon_high = self._fft_filter(recon, self.high_mask)\n        target_low = self._fft_filter(target, self.low_mask)\n        target_mid = self._fft_filter(target, self.mid_mask)\n        target_high = self._fft_filter(target, self.high_mask)\n\n        # L1 on low-freq (structure)\n        rec_loss_low = F.l1_loss(recon_low, target_low)\n\n        # LPIPS on mid-freq (perceptual)\n        p_loss = self.lpips_fn(recon_mid.float(), target_mid.float()).mean()\n\n        # Gradient matching on high-freq\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_high),\n            self._grad_magnitude(target_high)\n        )\n\n        # KL divergence\n        kl_loss = posterior.kl().mean()\n\n        # Adversarial losses\n        gen_loss, feat_loss, disc_loss = self._compute_gan_and_feat_loss(recon, target)\n\n        # Update discriminator\n        disc_active = step >= self.disc_start\n        if disc_active:\n            self.disc_opt.zero_grad()\n            disc_loss.backward()\n            torch.nn.utils.clip_grad_norm_(self.disc.parameters(), 1.0)\n            self.disc_opt.step()\n\n        # Adaptive weight calculation (taming-transformers style)\n        disc_factor = 1.0 if disc_active else 0.0\n\n        # Total loss\n        loss = (rec_loss_low +\n                self.perceptual_weight * p_loss +\n                self.grad_weight * grad_loss +\n                self.kl_weight * kl_loss +\n                disc_factor * (self.gan_weight * gen_loss +\n                               self.feat_match_weight * feat_loss))\n\n        return loss, {\n            \"rec_loss_low\": rec_loss_low.item(),\n            \"p_loss\": p_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"gen_loss\": gen_loss.item(),\n            \"feat_loss\": feat_loss.item(),\n            \"disc_loss\": disc_loss.item(),\n            \"disc_factor\": disc_factor,\n        }\n\n\ndef _build_disc_sn():\n    \"\"\"Build NLayer discriminator with spectral norm.\"\"\"\n    from torch.nn.utils import spectral_norm\n\n    class NLayerDiscWithFeatures(nn.Module):\n        def __init__(self):\n            super().__init__()\n            self.features = []\n            nf = 64\n            layers = []\n            # Input: [B, 3, 32, 32]\n            layers.append(spectral_norm(nn.Conv2d(3, nf, 4, stride=2, padding=1)))\n            layers.append(nn.LeakyReLU(0.2, True))\n            # [B, 64, 16, 16]\n            nf_prev = nf\n            nf = 128\n            layers.append(spectral_norm(nn.Conv2d(nf_prev, nf, 4, stride=2, padding=1)))\n            layers.append(nn.LeakyReLU(0.2, True))\n            # [B, 128, 8, 8]\n            nf_prev = nf\n            nf = 256\n            layers.append(spectral_norm(nn.Conv2d(nf_prev, nf, 4, stride=2, padding=1)))\n            layers.append(nn.LeakyReLU(0.2, True))\n            # [B, 256, 4, 4]\n            nf_prev = nf\n            nf = 512\n            layers.append(spectral_norm(nn.Conv2d(nf_prev, nf, 4, stride=1, padding=1)))\n            layers.append(nn.LeakyReLU(0.2, True))\n            # [B, 512, 4, 4]\n            layers.append(spectral_norm(nn.Conv2d(nf, 1, 4, stride=1, padding=1)))\n            self.model = nn.Sequential(*layers)\n\n        def forward(self, x, return_features=False):\n            feats = []\n            for layer in self.model:\n                x = layer(x)\n                if isinstance(layer, nn.Conv2d) and len(feats) < 4:\n                    feats.append(x)\n            output = x\n            if return_features:\n                return output, feats\n            return output\n\n    return NLayerDiscWithFeatures()\n\n\ndef _build_discriminator():\n    \"\"\"Placeholder for the _build_sn import.\"\"\"\n    class Dummy(nn.Module):\n        def forward(self, x): return x\n    return Dummy()",
5  "start_line": 32,
6  "end_line": 134
7}

1{
2  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Augmented Adversarial VAE Loss.\n\n    Combines standard VQGAN losses with frequency-domain regularization:\n    - L1 + LPIPS + KL (standard reconstruction losses)\n    - Adversarial hinge loss with spectral norm discriminator\n    - Feature matching loss\n    - Novel: frequency-domain high-pass gradient matching\n    - Novel: adaptive perceptual weighting via gradient sensitivity\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.perceptual_weight = 0.5\n\n        # Discriminator with spectral norm and feature extraction\n        self.disc = _NLayerDiscriminatorWithSN().to(device)\n        self.disc_opt = torch.optim.Adam(self.disc.parameters(), lr=1e-4, betas=(0.5, 0.9))\n        self.disc_start = 5000\n        self.feat_match_weight = 1.0\n\n        # Precompute high-pass frequency mask for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n        # High-pass: 1 / (1 + (r_cutoff / r)^(2n))\n        high_pass = 1.0 / (1.0 + (0.50 / r_norm) ** 4)\n        self.register_buffer('high_mask', high_pass.view(1, 1, H, W))\n\n    def _fft_highpass(self, img):\n        \"\"\"Extract high-frequency components.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * self.high_mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # Standard VQGAN losses\n        rec_loss = F.l1_loss(recon, target)\n        p_loss = self.lpips_fn(recon.float(), target.float()).mean()\n        kl_loss = posterior.kl().mean()\n\n        # Adversarial losses\n        fake_pred, fake_feats = self.disc(recon, return_features=True)\n        gen_loss = -fake_pred.mean()\n\n        real_pred, real_feats = self.disc(target, return_features=True)\n        fake_pred_det, _ = self.disc(recon.detach(), return_features=True)\n        disc_loss = F.relu(1.0 - real_pred).mean() + F.relu(1.0 + fake_pred_det).mean()\n\n        # Feature matching\n        feat_loss = 0.0\n        for rf, ff in zip(real_feats, fake_feats):\n            feat_loss += F.l1_loss(ff, rf)\n\n        # Update discriminator\n        disc_active = step >= self.disc_start\n        if disc_active:\n            self.disc_opt.zero_grad()\n            disc_loss.backward()\n            torch.nn.utils.clip_grad_norm_(self.disc.parameters(), 1.0)\n            self.disc_opt.step()\n\n        disc_factor = 1.0 if disc_active else 0.0\n\n        # Novel: frequency-domain gradient matching on high-freq components\n        recon_hp = self._fft_highpass(recon)\n        target_hp = self._fft_highpass(target)\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_hp),\n            self._grad_magnitude(target_hp)\n        )\n\n        # Total loss\n        loss = (rec_loss +\n                self.perceptual_weight * p_loss +\n                self.kl_weight * kl_loss +\n                disc_factor * (gen_loss + self.feat_match_weight * feat_loss) +\n                0.2 * grad_loss)\n\n        return loss, {\n            \"rec_loss\": rec_loss.item(),\n            \"p_loss\": p_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"gen_loss\": gen_loss.item(),\n            \"feat_loss\": feat_loss.item(),\n            \"disc_loss\": disc_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"disc_factor\": disc_factor,\n        }\n\n\nclass _NLayerDiscriminatorWithSN(nn.Module):\n    \"\"\"NLayer Discriminator with spectral normalization and feature extraction.\"\"\"\n    def __init__(self):\n        super().__init__()\n        from torch.nn.utils import spectral_norm as sn\n        nf = 64\n        self.convs = nn.ModuleList([\n            sn(nn.Conv2d(3, nf, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf, nf*2, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf*2, nf*4, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf*4, nf*4, 4, stride=1, padding=1)),\n        ])\n        self.act = nn.LeakyReLU(0.2, inplace=True)\n        self.conv_out = sn(nn.Conv2d(nf*4, 1, 4, stride=1, padding=1))\n\n    def forward(self, x, return_features=False):\n        feats = []\n        for conv in self.convs:\n            x = conv(x)\n            x = self.act(x)\n            feats.append(x)\n        out = self.conv_out(x)\n        if return_features:\n            return out, feats\n        return out",
3  "start_line": 32,
4  "end_line": 230,
5  "filename": "diffusers-main/custom_train.py"
6}

1{
2  "op": "replace",
3  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Augmented Adversarial VAE Loss.\n\n    Combines standard VQGAN losses with frequency-domain regularization:\n    - L1 + LPIPS + KL (standard reconstruction losses)\n    - Adversarial hinge loss with spectral norm discriminator\n    - Feature matching loss\n    - Novel: frequency-domain high-pass gradient matching for fine detail\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.perceptual_weight = 0.5\n\n        # Discriminator with spectral norm and feature extraction\n        self.disc = _NLayerDiscriminatorWithSN().to(device)\n        self.disc_opt = torch.optim.Adam(self.disc.parameters(), lr=1e-4, betas=(0.5, 0.9))\n        self.disc_start = 5000\n        self.feat_match_weight = 1.0\n\n        # Precompute high-pass frequency mask for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n        # High-pass: 1 / (1 + (r_cutoff / r)^(2n))\n        high_pass = 1.0 / (1.0 + (0.50 / r_norm) ** 4)\n        self.register_buffer('high_mask', high_pass.view(1, 1, H, W))\n\n    def _fft_highpass(self, img):\n        \"\"\"Extract high-frequency components.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * self.high_mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # Standard VQGAN losses\n        rec_loss = F.l1_loss(recon, target)\n        p_loss = self.lpips_fn(recon.float(), target.float()).mean()\n        kl_loss = posterior.kl().mean()\n\n        # Adversarial losses\n        fake_pred, fake_feats = self.disc(recon, return_features=True)\n        gen_loss = -fake_pred.mean()\n\n        real_pred, real_feats = self.disc(target, return_features=True)\n        fake_pred_det, _ = self.disc(recon.detach(), return_features=True)\n        disc_loss = F.relu(1.0 - real_pred).mean() + F.relu(1.0 + fake_pred_det).mean()\n\n        # Feature matching\n        feat_loss = 0.0\n        for rf, ff in zip(real_feats, fake_feats):\n            feat_loss += F.l1_loss(ff, rf)\n\n        # Update discriminator\n        disc_active = step >= self.disc_start\n        if disc_active:\n            self.disc_opt.zero_grad()\n            disc_loss.backward()\n            torch.nn.utils.clip_grad_norm_(self.disc.parameters(), 1.0)\n            self.disc_opt.step()\n\n        disc_factor = 1.0 if disc_active else 0.0\n\n        # Novel: frequency-domain gradient matching on high-freq components\n        recon_hp = self._fft_highpass(recon)\n        target_hp = self._fft_highpass(target)\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_hp),\n            self._grad_magnitude(target_hp)\n        )\n\n        # Total loss\n        loss = (rec_loss +\n                self.perceptual_weight * p_loss +\n                self.kl_weight * kl_loss +\n                disc_factor * (gen_loss + self.feat_match_weight * feat_loss) +\n                0.2 * grad_loss)\n\n        return loss, {\n            \"rec_loss\": rec_loss.item(),\n            \"p_loss\": p_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"gen_loss\": gen_loss.item(),\n            \"feat_loss\": feat_loss.item(),\n            \"disc_loss\": disc_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"disc_factor\": disc_factor,\n        }\n\n\nclass _NLayerDiscriminatorWithSN(nn.Module):\n    \"\"\"NLayer Discriminator with spectral normalization and feature extraction.\"\"\"\n    def __init__(self):\n        super().__init__()\n        from torch.nn.utils import spectral_norm as sn\n        nf = 64\n        self.convs = nn.ModuleList([\n            sn(nn.Conv2d(3, nf, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf, nf*2, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf*2, nf*4, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf*4, nf*4, 4, stride=1, padding=1)),\n        ])\n        self.act = nn.LeakyReLU(0.2, inplace=True)\n        self.conv_out = sn(nn.Conv2d(nf*4, 1, 4, stride=1, padding=1))\n\n    def forward(self, x, return_features=False):\n        feats = []\n        for conv in self.convs:\n            x = conv(x)\n            x = self.act(x)\n            feats.append(x)\n        out = self.conv_out(x)\n        if return_features:\n            return out, feats\n        return out",
4  "start_line": 32,
5  "end_line": 230,
6  "filename": "diffusers-main/custom_train.py"
7}

1{
2  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Decomposed Multi-Component VAE Loss with SSIM.\n\n    Novel components:\n    1. Frequency-band decomposition (Butterworth filters) for targeted\n       per-band loss: L1 on structure (low), LPIPS on perception (mid),\n       gradient matching on details (high).\n    2. SSIM loss for structural preservation at full resolution.\n    3. Progressive weighting that shifts emphasis from coarse to fine\n       structure as training progresses.\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n\n        # Precompute frequency band masks for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n\n        # Butterworth filters for low, mid, high bands\n        n_order = 2\n        low_cutoff = 0.22\n        mid_low = 0.22\n        mid_high = 0.50\n        # Low-pass\n        low_pass = 1.0 / (1.0 + (r_norm / low_cutoff) ** (2 * n_order))\n        # High-pass\n        high_pass = 1.0 / (1.0 + (mid_high / r_norm) ** (2 * n_order))\n        # Band-pass (mid)\n        mid_pass = (1.0 - low_pass) * high_pass\n\n        self.register_buffer('low_mask', low_pass.view(1, 1, H, W))\n        self.register_buffer('mid_mask', mid_pass.view(1, 1, H, W))\n        self.register_buffer('high_mask', (1.0 - high_pass).view(1, 1, H, W))\n\n        # Precompute Gaussian kernel for SSIM\n        self._ssim_kernel = self._make_gaussian_kernel(11, 1.5, device)\n\n    def _make_gaussian_kernel(self, size, sigma, device):\n        \"\"\"Create 2D Gaussian kernel.\"\"\"\n        coords = torch.arange(size, device=device).float() - (size - 1) / 2\n        g = torch.exp(-coords ** 2 / (2 * sigma ** 2))\n        g = g / g.sum()\n        kernel_1d = g.view(1, 1, size, 1)\n        kernel_2d = g.view(1, 1, 1, size)\n        return kernel_1d, kernel_2d\n\n    def _gaussian_blur(self, img):\n        \"\"\"Apply Gaussian blur using separable convolutions.\"\"\"\n        k1, k2 = self._ssim_kernel\n        b, c, h, w = img.shape\n        img = img.view(b * c, 1, h, w)\n        img = F.conv2d(img, k1, padding=(5, 0))\n        img = F.conv2d(img, k2, padding=(0, 5))\n        return img.view(b, c, h, w)\n\n    def _ssim_loss(self, recon, target):\n        \"\"\"Compute SSIM loss (1 - SSIM).\"\"\"\n        C1 = 0.01 ** 2\n        C2 = 0.03 ** 2\n\n        mu1 = self._gaussian_blur(recon)\n        mu2 = self._gaussian_blur(target)\n\n        mu1_sq = mu1 ** 2\n        mu2_sq = mu2 ** 2\n        mu1_mu2 = mu1 * mu2\n\n        sigma1_sq = self._gaussian_blur(recon ** 2) - mu1_sq\n        sigma2_sq = self._gaussian_blur(target ** 2) - mu2_sq\n        sigma12 = self._gaussian_blur(recon * target) - mu1_mu2\n\n        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / \\\n                   ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n        return (1.0 - ssim_map.mean()) / 2.0  # /2 to normalize to similar scale as L1\n\n    def _fft_filter(self, img, mask):\n        \"\"\"Apply frequency-domain filter.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # Frequency decomposition\n        recon_low = self._fft_filter(recon, self.low_mask)\n        recon_mid = self._fft_filter(recon, self.mid_mask)\n        recon_high = self._fft_filter(recon, self.high_mask)\n        target_low = self._fft_filter(target, self.low_mask)\n        target_mid = self._fft_filter(target, self.mid_mask)\n        target_high = self._fft_filter(target, self.high_mask)\n\n        # L1 on low-freq (structure)\n        rec_loss_low = F.l1_loss(recon_low, target_low)\n\n        # LPIPS on mid-freq (perceptual features)\n        p_loss = self.lpips_fn(recon_mid.float(), target_mid.float()).mean()\n\n        # Gradient matching on high-freq (edge sharpness)\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_high),\n            self._grad_magnitude(target_high)\n        )\n\n        # SSIM loss for structural preservation\n        ssim_loss = self._ssim_loss(recon, target)\n\n        # KL divergence\n        kl_loss = posterior.kl().mean()\n\n        # Fixed weights (progressive scheduling didn't help much)\n        loss = (1.0 * rec_loss_low +\n                0.5 * p_loss +\n                0.3 * grad_loss +\n                0.4 * ssim_loss +\n                self.kl_weight * kl_loss)\n\n        return loss, {\n            \"rec_loss_low\": rec_loss_low.item(),\n            \"p_loss\": p_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"ssim_loss\": ssim_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n        }",
3  "end_line": 163,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 32
7}

1{
2  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Decomposed Perceptual VAE Loss with SSIM.\n\n    Novel components:\n    1. Full-resolution L1 + LPIPS + KL (standard VQGAN-style reconstruction)\n    2. Frequency-band gradient matching: high-frequency components are\n       explicitly regularized via gradient magnitude matching in the\n       Fourier domain, preserving fine edges and texture.\n    3. Multi-scale SSIM loss computed at native resolution and\n       downsampled (16x16) resolutions for structural coherence.\n    4. Frequency-weighted L2 penalty on high-frequency residuals to\n       discourage blur without over-regularizing.\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.perceptual_weight = 0.5\n\n        # Precompute frequency band masks for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n\n        # High-pass Butterworth filter\n        n_order = 2\n        high_cutoff = 0.40\n        high_pass = 1.0 / (1.0 + (high_cutoff / r_norm) ** (2 * n_order))\n        self.register_buffer('high_mask', high_pass.view(1, 1, H, W))\n\n        # Precompute SSIM Gaussian kernels for 32x32\n        self._ssim_kernel_11 = self._make_gaussian_kernel(11, 1.5, device)\n        # Smaller kernel for downsampled SSIM\n        self._ssim_kernel_5 = self._make_gaussian_kernel(5, 0.8, device)\n\n    def _make_gaussian_kernel(self, size, sigma, device):\n        \"\"\"Create 2D Gaussian kernel.\"\"\"\n        coords = torch.arange(size, device=device).float() - (size - 1) / 2\n        g = torch.exp(-coords ** 2 / (2 * sigma ** 2))\n        g = g / g.sum()\n        k1 = g.view(1, 1, size, 1)\n        k2 = g.view(1, 1, 1, size)\n        return k1, k2\n\n    def _gaussian_blur(self, img, kernel_1d, kernel_2d):\n        \"\"\"Apply Gaussian blur using separable convolutions.\"\"\"\n        k1, k2 = kernel_1d, kernel_2d\n        b, c, h, w = img.shape\n        pad = (k1.shape[2] - 1) // 2\n        img = img.view(b * c, 1, h, w)\n        img = F.conv2d(img, k1, padding=(pad, 0))\n        img = F.conv2d(img, k2, padding=(0, pad))\n        return img.view(b, c, h, w)\n\n    def _ssim_loss(self, recon, target, kernel_1d, kernel_2d):\n        \"\"\"Compute SSIM loss (1 - SSIM).\"\"\"\n        C1 = 0.01 ** 2\n        C2 = 0.03 ** 2\n        mu1 = self._gaussian_blur(recon, kernel_1d, kernel_2d)\n        mu2 = self._gaussian_blur(target, kernel_1d, kernel_2d)\n        mu1_sq = mu1 ** 2\n        mu2_sq = mu2 ** 2\n        mu1_mu2 = mu1 * mu2\n        sigma1_sq = self._gaussian_blur(recon ** 2, kernel_1d, kernel_2d) - mu1_sq\n        sigma2_sq = self._gaussian_blur(target ** 2, kernel_1d, kernel_2d) - mu2_sq\n        sigma12 = self._gaussian_blur(recon * target, kernel_1d, kernel_2d) - mu1_mu2\n        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / \\\n                   ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n        return (1.0 - ssim_map.mean()) / 2.0\n\n    def _fft_highpass(self, img):\n        \"\"\"Extract high-frequency components.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * self.high_mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # Full-resolution reconstruction losses (like VQGAN)\n        rec_loss = F.l1_loss(recon, target)\n        p_loss = self.lpips_fn(recon.float(), target.float()).mean()\n\n        # Novel: Frequency-domain gradient matching on high-freq components\n        recon_hp = self._fft_highpass(recon)\n        target_hp = self._fft_highpass(target)\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_hp),\n            self._grad_magnitude(target_hp)\n        )\n\n        # Novel: Multi-scale SSIM (native + downsampled)\n        ssim_loss = self._ssim_loss(recon, target,\n                                     self._ssim_kernel_11[0], self._ssim_kernel_11[1])\n\n        # Downsampled SSIM for coarse structure\n        recon_ds = F.avg_pool2d(recon, 2)\n        target_ds = F.avg_pool2d(target, 2)\n        ssim_loss_ds = self._ssim_loss(recon_ds, target_ds,\n                                        self._ssim_kernel_5[0], self._ssim_kernel_5[1])\n\n        # KL divergence\n        kl_loss = posterior.kl().mean()\n\n        # Total loss\n        loss = (rec_loss +\n                self.perceptual_weight * p_loss +\n                0.15 * grad_loss +\n                0.3 * ssim_loss +\n                0.1 * ssim_loss_ds +\n                self.kl_weight * kl_loss)\n\n        return loss, {\n            \"rec_loss\": rec_loss.item(),\n            \"p_loss\": p_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"ssim_loss\": ssim_loss.item(),\n            \"ssim_loss_ds\": ssim_loss_ds.item(),\n            \"kl_loss\": kl_loss.item(),\n        }",
3  "end_line": 172,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 32
7}

1{
2  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Decomposed Perceptual VAE Loss with SSIM.\n\n    Novel components:\n    1. Full-resolution L1 + LPIPS + KL (standard VQGAN-style reconstruction)\n    2. Frequency-domain gradient matching: high-frequency components are\n       explicitly regularized via gradient magnitude matching in the\n       Fourier domain, preserving fine edges and texture.\n    3. Multi-scale SSIM loss computed at native resolution and\n       downsampled resolutions for structural coherence.\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.perceptual_weight = 0.5\n\n        # Precompute frequency band masks for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n\n        # High-pass Butterworth filter\n        n_order = 2\n        high_cutoff = 0.40\n        high_pass = 1.0 / (1.0 + (high_cutoff / r_norm) ** (2 * n_order))\n        self.register_buffer('high_mask', high_pass.view(1, 1, H, W))\n\n        # Precompute SSIM Gaussian kernels for 32x32\n        self._ssim_kernel_11 = self._make_gaussian_kernel(11, 1.5, device)\n        # Smaller kernel for downsampled SSIM\n        self._ssim_kernel_5 = self._make_gaussian_kernel(5, 0.8, device)\n\n    def _make_gaussian_kernel(self, size, sigma, device):\n        \"\"\"Create 2D Gaussian kernel.\"\"\"\n        coords = torch.arange(size, device=device).float() - (size - 1) / 2\n        g = torch.exp(-coords ** 2 / (2 * sigma ** 2))\n        g = g / g.sum()\n        k1 = g.view(1, 1, size, 1)\n        k2 = g.view(1, 1, 1, size)\n        return k1, k2\n\n    def _gaussian_blur(self, img, kernel_1d, kernel_2d):\n        \"\"\"Apply Gaussian blur using separable convolutions.\"\"\"\n        k1, k2 = kernel_1d, kernel_2d\n        b, c, h, w = img.shape\n        pad = (k1.shape[2] - 1) // 2\n        img = img.view(b * c, 1, h, w)\n        img = F.conv2d(img, k1, padding=(pad, 0))\n        img = F.conv2d(img, k2, padding=(0, pad))\n        return img.view(b, c, h, w)\n\n    def _ssim_loss(self, recon, target, kernel_1d, kernel_2d):\n        \"\"\"Compute SSIM loss (1 - SSIM).\"\"\"\n        C1 = 0.01 ** 2\n        C2 = 0.03 ** 2\n        mu1 = self._gaussian_blur(recon, kernel_1d, kernel_2d)\n        mu2 = self._gaussian_blur(target, kernel_1d, kernel_2d)\n        mu1_sq = mu1 ** 2\n        mu2_sq = mu2 ** 2\n        mu1_mu2 = mu1 * mu2\n        sigma1_sq = self._gaussian_blur(recon ** 2, kernel_1d, kernel_2d) - mu1_sq\n        sigma2_sq = self._gaussian_blur(target ** 2, kernel_1d, kernel_2d) - mu2_sq\n        sigma12 = self._gaussian_blur(recon * target, kernel_1d, kernel_2d) - mu1_mu2\n        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / \\\n                   ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))\n        return (1.0 - ssim_map.mean()) / 2.0\n\n    def _fft_highpass(self, img):\n        \"\"\"Extract high-frequency components.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * self.high_mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # Full-resolution reconstruction losses (like VQGAN)\n        rec_loss = F.l1_loss(recon, target)\n        p_loss = self.lpips_fn(recon.float(), target.float()).mean()\n\n        # Novel: Frequency-domain gradient matching on high-freq components\n        recon_hp = self._fft_highpass(recon)\n        target_hp = self._fft_highpass(target)\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_hp),\n            self._grad_magnitude(target_hp)\n        )\n\n        # Novel: Multi-scale SSIM (native + downsampled)\n        ssim_loss = self._ssim_loss(recon, target,\n                                     self._ssim_kernel_11[0], self._ssim_kernel_11[1])\n\n        # Downsampled SSIM for coarse structure\n        recon_ds = F.avg_pool2d(recon, 2)\n        target_ds = F.avg_pool2d(target, 2)\n        ssim_loss_ds = self._ssim_loss(recon_ds, target_ds,\n                                        self._ssim_kernel_5[0], self._ssim_kernel_5[1])\n\n        # KL divergence\n        kl_loss = posterior.kl().mean()\n\n        # Total loss\n        loss = (rec_loss +\n                self.perceptual_weight * p_loss +\n                0.15 * grad_loss +\n                0.3 * ssim_loss +\n                0.1 * ssim_loss_ds +\n                self.kl_weight * kl_loss)\n\n        return loss, {\n            \"rec_loss\": rec_loss.item(),\n            \"p_loss\": p_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"ssim_loss\": ssim_loss.item(),\n            \"ssim_loss_ds\": ssim_loss_ds.item(),\n            \"kl_loss\": kl_loss.item(),\n        }",
3  "end_line": 165,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 32
7}

1{
2  "content": "class VAELoss(nn.Module):\n    \"\"\"Frequency-Augmented VQGAN-style VAE Loss.\n\n    Combines the VQGAN adversarial framework with novel frequency-domain\n    gradient matching for enhanced detail preservation:\n    - L1 pixel reconstruction\n    - LPIPS perceptual similarity\n    - KL divergence regularization\n    - Adversarial hinge loss with spectral norm NLayer discriminator\n    - Feature matching loss from discriminator intermediate features\n    - Novel: Fourier-domain gradient matching on high-frequency band\n    \"\"\"\n\n    def __init__(self, device):\n        super().__init__()\n        import lpips\n        self.lpips_fn = lpips.LPIPS(net='vgg').to(device)\n        self.lpips_fn.eval()\n        for p in self.lpips_fn.parameters():\n            p.requires_grad_(False)\n        self.device = device\n        self.kl_weight = 1e-6\n        self.perceptual_weight = 0.5\n        self.feat_match_weight = 1.0\n\n        # Discriminator with spectral normalization\n        from torch.nn.utils import spectral_norm as sn\n        nf = 64\n        self.disc_convs = nn.ModuleList([\n            sn(nn.Conv2d(3, nf, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf, nf*2, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf*2, nf*4, 4, stride=2, padding=1)),\n            sn(nn.Conv2d(nf*4, nf*4, 4, stride=1, padding=1)),\n        ])\n        self.disc_act = nn.LeakyReLU(0.2, inplace=True)\n        self.disc_out = sn(nn.Conv2d(nf*4, 1, 4, stride=1, padding=1))\n        self.disc_opt = torch.optim.Adam(\n            list(self.disc_convs.parameters()) + list(self.disc_out.parameters()),\n            lr=1e-4, betas=(0.5, 0.9))\n        self.disc_start = 5000\n\n        # Precompute high-pass frequency mask for 32x32\n        H, W = 32, 32\n        cy, cx = H // 2, W // 2\n        y = torch.arange(H, device=device).view(H, 1).repeat(1, W)\n        x = torch.arange(W, device=device).view(1, W).repeat(H, 1)\n        r = torch.sqrt((x - cx) ** 2 + (y - cy) ** 2)\n        r_max = math.sqrt(cx ** 2 + cy ** 2)\n        r_norm = r / r_max\n        # High-pass Butterworth\n        high_cutoff = 0.40\n        high_pass = 1.0 / (1.0 + (high_cutoff / r_norm) ** 4)\n        self.register_buffer('high_mask', high_pass.view(1, 1, H, W))\n\n    def _disc_forward(self, x, return_features=False):\n        \"\"\"Forward pass through discriminator.\"\"\"\n        feats = []\n        for conv in self.disc_convs:\n            x = conv(x)\n            x = self.disc_act(x)\n            feats.append(x)\n        out = self.disc_out(x)\n        if return_features:\n            return out, feats\n        return out\n\n    def _fft_highpass(self, img):\n        \"\"\"Extract high-frequency components via FFT.\"\"\"\n        f = torch.fft.fft2(img, dim=(-2, -1))\n        f_filtered = f * self.high_mask\n        return torch.fft.ifft2(f_filtered, dim=(-2, -1)).real\n\n    def _grad_magnitude(self, img):\n        \"\"\"Compute gradient magnitude using finite differences.\"\"\"\n        dx = F.pad(img[:, :, :, 1:] - img[:, :, :, :-1], (0, 1, 0, 0))\n        dy = F.pad(img[:, :, 1:, :] - img[:, :, :-1, :], (0, 0, 0, 1))\n        return torch.sqrt(dx ** 2 + dy ** 2 + 1e-8)\n\n    def forward(self, recon, target, posterior, step):\n        # L1 pixel loss\n        rec_loss = F.l1_loss(recon, target)\n\n        # LPIPS perceptual loss\n        p_loss = self.lpips_fn(recon.float(), target.float()).mean()\n\n        # KL divergence\n        kl_loss = posterior.kl().mean()\n\n        # Adversarial losses\n        fake_pred, fake_feats = self._disc_forward(recon, return_features=True)\n        gen_loss = -fake_pred.mean()\n\n        real_pred, real_feats = self._disc_forward(target, return_features=True)\n        fake_pred_det, _ = self._disc_forward(recon.detach(), return_features=True)\n        disc_loss = F.relu(1.0 - real_pred).mean() + F.relu(1.0 + fake_pred_det).mean()\n\n        # Feature matching loss\n        feat_loss = 0.0\n        for rf, ff in zip(real_feats, fake_feats):\n            feat_loss += F.l1_loss(ff, rf)\n\n        # Update discriminator (separate from generator gradients)\n        disc_active = step >= self.disc_start\n        if disc_active:\n            self.disc_opt.zero_grad()\n            disc_loss.backward()\n            torch.nn.utils.clip_grad_norm_(\n                list(self.disc_convs.parameters()) + list(self.disc_out.parameters()), 1.0)\n            self.disc_opt.step()\n\n        disc_factor = 1.0 if disc_active else 0.0\n\n        # Novel: Frequency-domain gradient matching on high-freq components\n        recon_hp = self._fft_highpass(recon)\n        target_hp = self._fft_highpass(target)\n        grad_loss = F.l1_loss(\n            self._grad_magnitude(recon_hp),\n            self._grad_magnitude(target_hp)\n        )\n\n        # Total loss\n        loss = (rec_loss +\n                self.perceptual_weight * p_loss +\n                self.kl_weight * kl_loss +\n                disc_factor * (gen_loss + self.feat_match_weight * feat_loss) +\n                0.15 * grad_loss)\n\n        return loss, {\n            \"rec_loss\": rec_loss.item(),\n            \"p_loss\": p_loss.item(),\n            \"kl_loss\": kl_loss.item(),\n            \"gen_loss\": gen_loss.item(),\n            \"feat_loss\": feat_loss.item(),\n            \"disc_loss\": disc_loss.item(),\n            \"grad_loss\": grad_loss.item(),\n            \"disc_factor\": disc_factor,\n        }",
3  "end_line": 163,
4  "filename": "diffusers-main/custom_train.py",
5  "op": "replace",
6  "start_line": 32
7}