feat(dynacell): modernize pix2pix3d training recipe (StyleGAN2/R3GAN-style)

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/pix2pix3d_unetvit_fit_ddp_modernized.yml

@@ -0,0 +1,31 @@
+# pix2pix3d_unetvit fit overlay — modernized recipe, 4-GPU DDP variant.
+#
+# Sibling of pix2pix3d_unetvit_fit_modernized.yml that swaps in the
+# GAN-aware DDP topology (find_unused_parameters=True). All model knobs
+# match the single-GPU modernized overlay — DDP runs at the same per-rank
+# settings.
+#
+# See pix2pix3d_unetvit_fit_modernized.yml for rationale and plan reference.
+base:
+  - ../../../../../../recipes/models/pix2pix3d_unetvit.yml
+  - ../../../../../../recipes/trainer/fit.yml
+  - ../../../../../../recipes/topology/ddp_4gpu_gan.yml
+model:
+  init_args:
+    loss_type: nonsat
+    lambda_adv: 1.0
+    r1_gamma: 10.0
+    r2_gamma: 0.0
+    r1_every: 16
+    ema_kimg: 10.0
+    lecam_gamma: 0.0
+    lecam_decay: 0.9
+    use_ema_at_predict: true
+    lr_g: 2.0e-4
+    lr_d: 2.0e-4
+    schedule: WarmupCosine
+    warmup_steps: 8500
+    warmup_multiplier: 1e-3
+trainer:
+  precision: bf16-mixed
+  max_epochs: 20

applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/pix2pix3d_unetvit_fit_modernized.yml

@@ -0,0 +1,51 @@
+# pix2pix3d_unetvit fit overlay — STYLEGAN2/R3GAN-STYLE MODERNIZED RECIPE.
+#
+# Opt-in modernization of the LSGAN baseline. See PR
+# https://github.com/mehta-lab/VisCy/pull/428 for the recipe and rationale.
+#
+# Differences from pix2pix3d_unetvit_fit.yml:
+#   - loss_type:   lsgan         -> nonsat (StyleGAN2 default; softplus-based)
+#   - lr_g/lr_d:   3e-4 / 3e-5   -> 2e-4 / 2e-4 (equal LR; matches StyleGAN2-ADA /
+#                                  StyleGAN3 / StyleGAN-XL / R3GAN; the 10x
+#                                  D-slower TTUR was a pre-R1-era hack)
+#   - r1_gamma:    0             -> 10.0  (Mescheder zero-centered grad penalty,
+#                                  StyleGAN2 256² default; ablate {1,5,10,20})
+#   - r1_every:    n/a           -> 16    (StyleGAN2 D_reg_interval; lazy reg
+#                                  with `* r1_every` unbiased rescaling)
+#   - ema_kimg:    null          -> 10.0  (G EMA half-life 10k images;
+#                                  StyleGAN2 256² default. Decay derived
+#                                  per-step from global batch size.)
+#   - lambda_adv:  implicit 1.0  -> 1.0   (now an explicit knob)
+#
+# Unchanged from the LSGAN baseline:
+#   - generator backbone (UNetViT3D)
+#   - discriminator backbone (MultiScalePatchGAN3D, num_scales=2, SN on)
+#   - lambda_l1 (=100.0 from recipes/models/pix2pix3d_unetvit.yml)
+#   - warmup_steps, warmup_multiplier, schedule, precision, max_epochs
+#
+# Architectural fairness with CellDiff / FNet3D / etc. is preserved: only
+# the training RECIPE differs from LSGAN; the model under test is still
+# pix2pix3d_unetvit.
+base:
+  - ../../../../../../recipes/models/pix2pix3d_unetvit.yml
+  - ../../../../../../recipes/trainer/fit.yml
+  - ../../../../../../recipes/topology/single_gpu.yml
+model:
+  init_args:
+    loss_type: nonsat
+    lambda_adv: 1.0
+    r1_gamma: 10.0
+    r2_gamma: 0.0
+    r1_every: 16
+    ema_kimg: 10.0
+    lecam_gamma: 0.0
+    lecam_decay: 0.9
+    use_ema_at_predict: true
+    lr_g: 2.0e-4
+    lr_d: 2.0e-4
+    schedule: WarmupCosine
+    warmup_steps: 8500
+    warmup_multiplier: 1e-3
+trainer:
+  precision: bf16-mixed
+  max_epochs: 20

applications/dynacell/configs/benchmarks/virtual_staining/er/pix2pix3d_unetvit/ipsc_confocal/train_4gpu_modernized.yml

@@ -0,0 +1,51 @@
+# pix2pix3d_unetvit MODERNIZED fit on ER (SEC61B) — 4-GPU DDP production leaf.
+#
+# Run A of the modernization recipe (see PR
+# https://github.com/mehta-lab/VisCy/pull/428).
+#
+# Sibling of train_4gpu.yml. Same experiment_id and W&B/save-dir structure,
+# but composes the MODERNIZED model overlay (nonsat + R1 + EMA + equal LR)
+# instead of the LSGAN baseline overlay. Checkpoints land in a distinct
+# subtree so they cannot collide with the in-flight LSGAN run.
+#
+# W&B run name and slurm job name carry `_modernized` to make the recipe
+# obvious in the dashboard.
+base:
+  - ../../../_internal/shared/model/train_sets/ipsc_confocal.yml
+  - ../../../_internal/shared/model/targets/er_sec61b.yml
+  - ../../../_internal/shared/model/data_overlays/unetvit3d_fit.yml
+  - ../../../_internal/shared/model/model_overlays/pix2pix3d_unetvit_fit_ddp_modernized.yml
+  - ../../../_internal/shared/model/launcher_profiles/mode_fit.yml
+  - ../../../_internal/shared/model/launcher_profiles/hardware_4gpu.yml
+  - ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml
+
+benchmark:
+  task: virtual_staining
+  organelle: er
+  train_set: ipsc_confocal
+  model_name: pix2pix3d_unetvit
+  experiment_id: er__ipsc_confocal__pix2pix3d_unetvit_modernized
+
+trainer:
+  logger:
+    init_args:
+      name: pix2pix3d_unetvit_iPSC_SEC61B_4gpu_modernized
+      save_dir: /hpc/projects/comp.micro/virtual_staining/models/cell_diff_vs_viscy/ipsc/sec61b/pix2pix3d_unetvit_modernized
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: step
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        # Modernized recipe monitors the EMA generator's val loss — that's
+        # the inference path (predict_step uses _inference_generator which
+        # prefers EMA when available).
+        monitor: loss/validate_ema
+        every_n_epochs: 1
+        save_top_k: 4
+        save_last: true
+        dirpath: /hpc/projects/comp.micro/virtual_staining/models/cell_diff_vs_viscy/ipsc/sec61b/pix2pix3d_unetvit_modernized/checkpoints
+
+launcher:
+  job_name: pix2pix3d_unetvit_SEC61B_4gpu_modernized
+  run_root: /hpc/projects/comp.micro/virtual_staining/models/cell_diff_vs_viscy/ipsc/sec61b/pix2pix3d_unetvit_modernized