fix smoke bugs + anomaly guards; runbook update

RUNBOOK_A100.md

@@ -49,3 +49,16 @@ bash scripts/run_stage3.sh   # S3 域微调 (20k step)
 - 训练脚本依赖 torchvision(官方 bsr/degradations.py import)。A100(Linux+torchvision)正常;
   本机 Windows 沙箱若 torchvision 报 `operator torchvision::nms does not exist`, 请在自己终端(非沙箱)运行。
 - 数据清单中的真实退化对仅 RealSR 400/99; 其余母本在 data/clean_all, data/clean。
+
+## 本地冒烟结果(2026-09-06, conda tvtest CPU)
+- conda 环境 torch 2.5.1 + torchvision 0.20.1 可用(需 PYTHONNOUSERSITE=1 屏蔽用户目录污染)
+- 修复清单(已同步到本仓库):
+  * common.py: LoRAConv2d cin/cout 未定义; LoRA 仅注入 Linear/1x1Conv(3x3 分解不匹配);
+    整体冻结仅训 LoRA; assemble_full_student 展开 up_blocks(ModuleList 不能进 Sequential);
+    load_diffusers_sd 自动 variant=fp16; 新增 is_finite/check_tensor/clip_and_check_grads/EMA/preview_grid
+  * train_lora.py/train_stage2.py: NaN/Inf 检测(输入/输出/损失/梯度)、梯度裁剪(clip_grad 1.0)、
+    skip-step 连续异常>20 中止、EMA(默认0.999)、周期预览图(vis_every)、num_workers 冒烟用0
+  * official/bsr/utils/img_process_util.py: filter2D .view->.reshape(非连续张量)
+  * configs/config_smoke.yml: gt_size 128->512(Net 输入需 LR128=HR512/4)
+- 冒烟结果: 退化管线 OK(HR512->LR128, 数值正常, 预览 logs/smoke_degrade_preview.png);
+  S1 LoRA 2步 OK(checkpoint/lora_ema/预览); S2 本地 CPU 内存不足被系统终止 -> 在 A100 跑 S2/GDPO 冒烟

configs/config_smoke.yml

@@ -1,37 +1,37 @@
-dataroot_gt: data/patches
-gt_size: 128
-iter_num: 2
-# Real-ESRGAN style degradation (shared by all stages)
-scale: 4
-resize_prob: [0.2, 0.7, 0.1]
-resize_range: [0.3, 1.5]
-gaussian_noise_prob: 0.5
-noise_range: [1, 15]
-poisson_scale_range: [0.05, 2.0]
-gray_noise_prob: 0.4
-jpeg_range: [60, 95]
-second_blur_prob: 0.5
-resize_prob2: [0.3, 0.4, 0.3]
-resize_range2: [0.6, 1.2]
-gaussian_noise_prob2: 0.5
-noise_range2: [1, 12]
-poisson_scale_range2: [0.05, 1.0]
-gray_noise_prob2: 0.4
-jpeg_range2: [60, 100]
-blur_kernel_size: 21
-kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
-kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
-sinc_prob: 0.1
-blur_sigma: [0.2, 1.5]
-betag_range: [0.5, 2.0]
-betap_range: [1, 1.5]
-blur_kernel_size2: 11
-kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
-kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
-sinc_prob2: 0.1
-blur_sigma2: [0.2, 1.0]
-betag_range2: [0.5, 2.0]
-betap_range2: [1, 1.5]
-final_sinc_prob: 0.8
-use_hflip: True
+dataroot_gt: data/patches
+gt_size: 512
+iter_num: 2
+# Real-ESRGAN style degradation (shared by all stages)
+scale: 4
+resize_prob: [0.2, 0.7, 0.1]
+resize_range: [0.3, 1.5]
+gaussian_noise_prob: 0.5
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2.0]
+gray_noise_prob: 0.4
+jpeg_range: [60, 95]
+second_blur_prob: 0.5
+resize_prob2: [0.3, 0.4, 0.3]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.5
+noise_range2: [1, 12]
+poisson_scale_range2: [0.05, 1.0]
+gray_noise_prob2: 0.4
+jpeg_range2: [60, 100]
+blur_kernel_size: 21
+kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob: 0.1
+blur_sigma: [0.2, 1.5]
+betag_range: [0.5, 2.0]
+betap_range: [1, 1.5]
+blur_kernel_size2: 11
+kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+sinc_prob2: 0.1
+blur_sigma2: [0.2, 1.0]
+betag_range2: [0.5, 2.0]
+betap_range2: [1, 1.5]
+final_sinc_prob: 0.8
+use_hflip: True
 use_rot: False

official/bsr/utils/img_process_util.py

@@ -22,13 +22,13 @@ def filter2D(img, kernel):
 
     if kernel.size(0) == 1:
         # apply the same kernel to all batch images
-        img = img.view(b * c, 1, ph, pw)
-        kernel = kernel.view(1, 1, k, k)
-        return F.conv2d(img, kernel, padding=0).view(b, c, h, w)
+        img = img.reshape(b * c, 1, ph, pw)
+        kernel = kernel.reshape(1, 1, k, k)
+        return F.conv2d(img, kernel, padding=0).reshape(b, c, h, w)
     else:
-        img = img.view(1, b * c, ph, pw)
-        kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k)
-        return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w)
+        img = img.reshape(1, b * c, ph, pw)
+        kernel = kernel.reshape(b, 1, k, k).repeat(1, c, 1, 1).reshape(b * c, 1, k, k)
+        return F.conv2d(img, kernel, groups=b * c).reshape(b, c, h, w)
 
 
 def usm_sharp(img, weight=0.5, radius=50, threshold=10):

src/common.py

@@ -2,7 +2,7 @@
 """AdcSR 工程公共工具：路径、模型装配、手工 LoRA 注入、教师加载、GDPO probe。
 训练脚本统一从这里 import，禁止各自重复实现装配逻辑。
 """
-import os, sys, copy, json, types
+import os, sys, copy, json, types, math
 from pathlib import Path
 
 REPO = Path(__file__).resolve().parents[1]
@@ -21,9 +21,15 @@ import torch.nn.functional as F
 # ---------------------------------------------------------------------------
 # 模型装配（与 official/test.py 全链一致）
 # ---------------------------------------------------------------------------
-def load_diffusers_sd(model_id, dtype=torch.float32, device="cpu"):
+def load_diffusers_sd(model_id, dtype=torch.float32, device="cpu", variant=None):
     from diffusers import StableDiffusionPipeline
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype).to(device)
+    if variant is None:
+        # ???? fp16 ???????; ??? variant="" ???
+        import os as _os
+        if _os.path.isdir(model_id) and _os.path.exists(_os.path.join(model_id, "unet", "diffusion_pytorch_model.fp16.safetensors")):
+            variant = "fp16"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype,
+                                                   variant=variant).to(device)
     return pipe.vae, pipe.unet, pipe.text_encoder, pipe.tokenizer
 
 def load_pruned_decoder(half_decoder_ckpt, device="cpu", dtype=torch.float32):
@@ -51,7 +57,7 @@ def assemble_full_student(unet, decoder, net_weights=None, device="cuda", dtype=
             sd = {k.replace("module.", "", 1): v for k, v in sd.items()}
         net.load_state_dict(sd, strict=True)
     net.to(device=device, dtype=dtype)
-    tail_mods = [decoder.up_blocks, decoder.conv_norm_out, decoder.conv_act, decoder.conv_out]
+    tail_mods = [*decoder.up_blocks, decoder.conv_norm_out, decoder.conv_act, decoder.conv_out]
     for m in tail_mods:
         m.to(device=device, dtype=dtype)
     full = nn.Sequential(net, *tail_mods)
@@ -135,9 +141,10 @@ class LoRAConv2d(nn.Module):
         self.conv = conv
         self.r = max(1, r)
         self.alpha = alpha
-        cin, cout = conv.in_channels, conv.out_channels
-        self.lora_a = nn.Parameter(torch.zeros(cin, self.r))
-        self.lora_b = nn.Parameter(torch.zeros(self.r, cout))
+        self.cin = conv.in_channels
+        self.cout = conv.out_channels
+        self.lora_a = nn.Parameter(torch.zeros(self.cin, self.r))
+        self.lora_b = nn.Parameter(torch.zeros(self.r, self.cout))
         nn.init.kaiming_uniform_(self.lora_a, a=5 ** 0.5)
         nn.init.zeros_(self.lora_b)
         self.requires_grad_(False)
@@ -148,8 +155,8 @@ class LoRAConv2d(nn.Module):
         y = self.conv(x)
         if self.training or True:
             # 1x1 conv low-rank: 输入cin->r->cout, 保持空间尺寸
-            z = F.conv2d(x, self.lora_a.t().view(self.r, cin, 1, 1))
-            z = F.conv2d(z, self.lora_b.t().view(cout, self.r, 1, 1))
+            z = F.conv2d(x, self.lora_a.t().view(self.r, self.cin, 1, 1))
+            z = F.conv2d(z, self.lora_b.t().view(self.cout, self.r, 1, 1))
             return y + self.alpha * z
         return y
 
@@ -186,10 +193,18 @@ def inject_lora(model, rank=64, alpha=1.0, skip_bias_norm=True, include=("conv",
         if include is not None and not any(s in n for s in include):
             continue
         parent, attr = _find_parent(model, n)
-        if isinstance(m, nn.Conv2d):
+        if isinstance(m, nn.Conv2d) and m.kernel_size == (1, 1):
             setattr(parent, attr, LoRAConv2d(m, rank, alpha))
+        elif isinstance(m, nn.Conv2d):
+            continue  # 3x3/stride>1 conv: 1x1 ???????, ??
         elif isinstance(m, nn.Linear):
             setattr(parent, attr, LoRALinear(m, rank, alpha))
+    # ????, ??? LoRA A/B ???(?????, ??"? LoRA ??")
+    model.requires_grad_(False)
+    for m in model.modules():
+        if isinstance(m, (LoRAConv2d, LoRALinear)):
+            m.lora_a.requires_grad_(True)
+            m.lora_b.requires_grad_(True)
     return model
 
 def _find_parent(model, name):
@@ -204,6 +219,86 @@ def lora_params(model):
         if p.requires_grad:
             yield p
 
+# ---------------------------------------------------------------------------
+# ???????(????/????/EMA/???) 2026-09-06
+# ---------------------------------------------------------------------------
+def is_finite(x):
+    """??/???????(? NaN/Inf)?"""
+    try:
+        if torch.is_tensor(x):
+            return bool(torch.isfinite(x.float()).all().item())
+        return bool(math.isfinite(float(x)))
+    except Exception:
+        return False
+
+def check_tensor(x, name, log=None):
+    """??/??????: ?? True=???"""
+    if x is None:
+        return False
+    if torch.is_tensor(x) and not is_finite(x):
+        msg = f"[anomaly] {name} contains NaN/Inf"
+        print(msg, flush=True)
+        if log is not None:
+            log(msg)
+        return True
+    return False
+
+def clip_and_check_grads(params, max_norm, log=None):
+    """???? + NaN/Inf ??; ?? True=????(??? step)?"""
+    grads = [p.grad for p in params if p.grad is not None]
+    bad = False
+    for g in grads:
+        if not is_finite(g):
+            bad = True
+            msg = "[anomaly] grad contains NaN/Inf; skip this optimizer step"
+            print(msg, flush=True)
+            if log is not None:
+                log(msg)
+            break
+    if bad:
+        return True
+    if max_norm and max_norm > 0 and grads:
+        total = torch.nn.utils.clip_grad_norm_(params, max_norm=max_norm)
+        if not is_finite(total):
+            msg = "[anomaly] grad total norm NaN; skip step"
+            print(msg, flush=True)
+            if log is not None:
+                log(msg)
+            return True
+    return False
+
+class EMA:
+    """??????(?? trainable/lora ??)?"""
+    def __init__(self, params, decay=0.999):
+        self.decay = decay
+        self.shadow = {id(p): p.detach().clone().float() for p in params if p.requires_grad}
+    @torch.no_grad()
+    def update(self, params):
+        d = self.decay
+        for p in params:
+            if not p.requires_grad or id(p) not in self.shadow:
+                continue
+            self.shadow[id(p)].mul_(d).add_(p.detach().float(), alpha=1 - d)
+    def state_dict(self, params):
+        return {id(p): self.shadow[id(p)] for p in params if id(p) in self.shadow}
+
+def preview_grid(tensors, path, vmin=-1.0, vmax=1.0):
+    """? [B,C,H,W] ??([-1,1]) ?????? PNG, ????????/?????"""
+    import numpy as np
+    from PIL import Image
+    ims = []
+    for t in tensors:
+        t = t.detach().float().clamp(vmin, vmax)
+        t = (t - vmin) / (vmax - vmin)
+        b = t[0].clamp(0, 1).permute(1, 2, 0).cpu().numpy()
+        ims.append(Image.fromarray((b * 255).astype(np.uint8)))
+    w = sum(im.width for im in ims); h = max(im.height for im in ims)
+    canvas = Image.new("RGB", (w, h), (0, 0, 0))
+    x = 0
+    for im in ims:
+        canvas.paste(im, (x, 0)); x += im.width
+    canvas.save(path, quality=92)
+
 def count_params(model, only_trainable=False):
     if only_trainable:
         return sum(p.numel() for p in model.parameters() if p.requires_grad)

src/train_lora.py

@@ -21,7 +21,8 @@ from omegaconf import OmegaConf
 
 from common import (ensure_official, load_diffusers_sd, load_pruned_decoder,
                     assemble_full_student, inject_lora, lora_params,
-                    count_params, build_net)
+                    count_params, build_net, is_finite, check_tensor,
+                    clip_and_check_grads, EMA, preview_grid)
 ensure_official()
 from dataset import RealESRGANDataset, RealESRGANDegrader  # official
 
@@ -127,6 +128,9 @@ def main():
     ap.add_argument("--no_bf16", dest="bf16", action="store_false")
     ap.add_argument("--seed", type=int, default=123)
     ap.add_argument("--num_workers", type=int, default=8)
+    ap.add_argument("--clip_grad", type=float, default=1.0, help="??????; 0=??")
+    ap.add_argument("--ema_decay", type=float, default=0.999, help="EMA ??; 0=??")
+    ap.add_argument("--vis_every", type=int, default=500, help="? N ?? LR/HR/?????")
     args = ap.parse_args()
     random.seed(args.seed); torch.manual_seed(args.seed)
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -175,9 +179,11 @@ def main():
         except Exception as e:
             log(f"[warn] pyiqa dists 不可用: {e}; 该损失置 0")
 
-    # ---- train loop ----
+    # ---- train loop (????/????/EMA/???) ----
     syn_iter = iter(syn_dl); real_iter = iter(real_dl) if real_dl else None
-    step = 0
+    step = 0; skip_streak = 0
+    ema = EMA(params, args.ema_decay) if args.ema_decay > 0 else None
+    name_of = {id(p): n for n, p in full.named_parameters() if p.requires_grad}
     full.train()
     optimizer.zero_grad(set_to_none=True)
     while step < args.steps:
@@ -193,6 +199,11 @@ def main():
             real_iter = iter(real_dl) if real_dl else None
             continue
         lr_t, hr_t = lr_t.to(device), hr_t.to(device)
+        if check_tensor(lr_t, "lr", log) or check_tensor(hr_t, "hr", log):
+            skip_streak += 1
+            if skip_streak > 20:
+                log("[anomaly] too many bad batches, abort"); break
+            continue
         if dtype == torch.bfloat16:
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 out = full(lr_t)
@@ -200,21 +211,50 @@ def main():
         else:
             out = full(lr_t)
             loss, items = _losses(out, hr_t, full, args, lpips_fn, dists_fn)
+        if check_tensor(out, "output", log):
+            skip_streak += 1
+            if skip_streak > 20:
+                log("[anomaly] too many bad outputs, abort"); break
+            optimizer.zero_grad(set_to_none=True)
+            continue
+        if not is_finite(loss):
+            log(f"[anomaly] loss NaN/Inf at step {step+1}; skip step")
+            skip_streak += 1
+            optimizer.zero_grad(set_to_none=True)
+            if skip_streak > 20:
+                log("[anomaly] too many bad losses, abort"); break
+            continue
+        skip_streak = 0
         (loss / args.grad_accum).backward()
         if (step + 1) % args.grad_accum == 0:
-            if scaler is not None:
-                scaler.step(optimizer); scaler.update()
+            if clip_and_check_grads(params, args.clip_grad, log):
+                optimizer.zero_grad(set_to_none=True)
             else:
-                optimizer.step()
-            optimizer.zero_grad(set_to_none=True)
-            sched.step()
+                if scaler is not None:
+                    scaler.step(optimizer); scaler.update()
+                else:
+                    optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+                sched.step()
+                if ema is not None:
+                    ema.update(params)
         if (step + 1) % 50 == 0:
             log(f"step {step+1}/{args.steps} loss {loss.item():.4f} " +
                 " ".join(f"{k}:{v:.4f}" for k, v in items.items()))
+        if args.vis_every > 0 and (step + 1) % args.vis_every == 0:
+            try:
+                preview_grid([lr_t.float()[:1], out.float()[:1], hr_t.float()[:1]],
+                             os.path.join(args.log_dir, f"step_{step+1:06d}.png"))
+            except Exception as e:
+                log(f"[warn] preview fail: {e}")
         if (step + 1) % args.save_every == 0:
             _save(full, args.out, step + 1)
+            if ema is not None:
+                _save_ema(ema, name_of, args.out, step + 1)
         step += 1
     _save(full, args.out, step)
+    if ema is not None:
+        _save_ema(ema, name_of, args.out, step)
     log("S1 done")
 
 def _losses(out, hr, full, args, lpips_fn, dists_fn):
@@ -250,5 +290,14 @@ def _save(full, out_dir, step):
     torch.save(net.state_dict(), os.path.join(out_dir, f"net_params_{step}.pkl"))
     torch.save(full.state_dict(), os.path.join(out_dir, f"full_params_{step}.pkl"))
 
+def _save_ema(ema, name_of, out_dir, step):
+    sd = {}
+    for pid, val in ema.shadow.items():
+        nm = name_of.get(pid)
+        if nm:
+            sd[nm] = val.detach().cpu().clone()
+    if sd:
+        torch.save(sd, os.path.join(out_dir, f"lora_ema_{step}.pkl"))
+
 if __name__ == "__main__":
     main()

src/train_stage2.py

@@ -17,7 +17,8 @@ from omegaconf import OmegaConf
 
 from common import (ensure_official, load_diffusers_sd, load_pruned_decoder,
                     build_net, build_discriminator_unet, load_osediff_teacher,
-                    load_gdpo_teacher, count_params)
+                    load_gdpo_teacher, count_params, is_finite, check_tensor,
+                    clip_and_check_grads, preview_grid)
 ensure_official()
 from dataset import RealESRGANDataset, RealESRGANDegrader  # official
 from train_lora import highfreq_mask  # 复用小波高频 mask
@@ -44,6 +45,8 @@ def main():
     ap.add_argument("--skip_ram", action="store_true")
     ap.add_argument("--bf16", action="store_true", default=True); ap.add_argument("--no_bf16", dest="bf16", action="store_false")
     ap.add_argument("--seed", type=int, default=123); ap.add_argument("--num_workers", type=int, default=8)
+    ap.add_argument("--clip_grad", type=float, default=1.0, help="??/?????????; 0=??")
+    ap.add_argument("--vis_every", type=int, default=500, help="? N ?? ??RGB/HR ??")
     args = ap.parse_args()
     random.seed(args.seed); torch.manual_seed(args.seed)
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -147,7 +150,9 @@ def main():
             for p in dists_fn.parameters(): p.requires_grad_(False)
         except Exception as e: log(f"[warn] pyiqa dists 不可用 {e}")
 
-    dl_iter = iter(dataloader); step = 0
+    dl_iter = iter(dataloader); step = 0; skip_streak = 0
+    params_s = list(net.parameters()) + list(tail.parameters())
+    params_D = [p for p in unet_D.parameters() if p.requires_grad]
     opt_s.zero_grad(set_to_none=True); opt_D.zero_grad(set_to_none=True)
     loss_D = torch.tensor(0.0, device=device)
     while step < args.steps:
@@ -157,6 +162,11 @@ def main():
             dl_iter = iter(dataloader); continue
         LR, HR = degrader.degrade(batch)          # [0,1] range
         LR, HR = LR.to(device), HR.to(device)
+        if check_tensor(LR, "lr", log) or check_tensor(HR, "hr", log):
+            skip_streak += 1
+            if skip_streak > 20:
+                log("[anomaly] too many bad batches, abort"); break
+            continue
         with torch.no_grad(), amp:
             tag = DAPE.generate_tag(ram_tf(LR))[0]
             text_input = tokenizer(tag, max_length=tokenizer.model_max_length,
@@ -174,6 +184,11 @@ def main():
             z0_t = decoder.conv_in(z0_t); z0_t = decoder.mid_block(z0_t)
             z0_gt = vae.post_quant_conv(HR_lat)
             z0_gt = decoder.conv_in(z0_gt); z0_gt = decoder.mid_block(z0_gt)
+        if check_tensor(z0_t, "z0_teacher", log) or check_tensor(z0_gt, "z0_gt", log):
+            skip_streak += 1
+            if skip_streak > 20:
+                log("[anomaly] too many bad teachers, abort"); break
+            continue
         with amp:
             z0_s = net(LRn)
             tB = torch.full((z0_s.shape[0],), 999, dtype=torch.long, device=device)
@@ -185,6 +200,7 @@ def main():
                 loss_distil = (m * (z0_s - z0_t).abs()).mean()
                 loss_adv = (m * F.softplus(-d_s)).mean()
             loss = args.w_distil * loss_distil + args.w_adv * loss_adv
+            img_s = None
             if args.w_lpips + args.w_dists + args.w_color > 0:
                 img_s = tail(z0_s)
                 if lpips_fn is not None and args.w_lpips > 0:
@@ -196,19 +212,39 @@ def main():
                     loss = loss + args.w_color * (
                         (img_s.float().mean(dim=(2, 3)) - HRn.float().mean(dim=(2, 3))).abs().mean()
                         + (img_s.float().std(dim=(2, 3)) - HRn.float().std(dim=(2, 3))).abs().mean())
+        if check_tensor(z0_s, "z0_student", log) or not is_finite(loss):
+            log(f"[anomaly] student/loss NaN at step {step+1}; skip step")
+            skip_streak += 1
+            opt_s.zero_grad(set_to_none=True); opt_D.zero_grad(set_to_none=True)
+            if skip_streak > 20:
+                log("[anomaly] too many bad steps, abort"); break
+            continue
+        skip_streak = 0
         (loss / args.grad_accum).backward()
         if (step + 1) % args.grad_accum == 0:
-            opt_s.step(); opt_s.zero_grad(set_to_none=True); sched.step()
+            if clip_and_check_grads(params_s, args.clip_grad, log):
+                opt_s.zero_grad(set_to_none=True)
+            else:
+                opt_s.step(); opt_s.zero_grad(set_to_none=True); sched.step()
             with amp:
                 tB = torch.full((z0_s.shape[0],), 999, dtype=torch.long, device=device)
                 pred_real = unet_D(z0_gt.detach(), tB, encoder_hidden_states=enc, return_dict=False)[0]
                 pred_fake = unet_D(z0_s.detach(), tB, encoder_hidden_states=enc, return_dict=False)[0]
                 loss_D = F.softplus(pred_fake).mean() + F.softplus(-pred_real).mean()
             (loss_D / args.grad_accum).backward()
-            opt_D.step(); opt_D.zero_grad(set_to_none=True)
+            if clip_and_check_grads(params_D, args.clip_grad, log):
+                opt_D.zero_grad(set_to_none=True)
+            else:
+                opt_D.step(); opt_D.zero_grad(set_to_none=True)
         if (step + 1) % 50 == 0:
             log(f"step {step+1}/{args.steps} g {loss.item():.4f} distil {loss_distil.item():.4f} "
                 f"adv {loss_adv.item():.4f} D {loss_D.item():.4f}")
+        if args.vis_every > 0 and (step + 1) % args.vis_every == 0 and img_s is not None:
+            try:
+                preview_grid([LRn.float()[:1], img_s.float()[:1], HRn.float()[:1]],
+                             os.path.join(args.log_dir, f"step_{step+1:06d}.png"))
+            except Exception as e:
+                log(f"[warn] preview fail: {e}")
         if (step + 1) % args.save_every == 0:
             torch.save(net.state_dict(), os.path.join(args.out, f"net_params_{step+1}.pkl"))
             torch.save(tail.state_dict(), os.path.join(args.out, f"tail_params_{step+1}.pkl"))