Go to ICLR 2026 Conference homepageMagicStain: High-Fidelity Pathology Image Virtual Staining via Guided Single-Step Diffusion
Keywords: Medical and Biological Imaging, Diffusion Models for Vision
Abstract: Virtual staining, which leverages computational methods to generate different stain styles from a pathology image of an already stained tissue section, offers a cost-effective alternative to chemical multiple staining. Despite extensive research based on generative adversarial networks (GANs) and diffusion models, achieving high-fidelity, high-quality, and computationally efficient results remains a significant challenge for virtual staining methods. While diffusion-based approaches typically produce more photorealistic images than GAN-based counterparts through multi-step sampling, this comes at the cost of high computational overhead and inference latency. In this paper, we propose MagicStain, a novel single-step diffusion model tailored for generating high-resolution virtual stains. Specifically, we adapt a pretrained single-step diffusion model to enable efficient virtual staining. By introducing pathology priors from a pretrained pathology-specific vision language model and integrating pathology- and structure-consistency losses on both the original images and the H-channel, MagicStain achieves high-fidelity and high-quality generations. To address the limitations of single-step diffusion models in high-resolution virtual staining, we further propose a two-stage progressive training strategy that enables high-resolution adaptation with low training cost. Extensive experiments on three virtual staining datasets, each involving translation between different staining dyes/biomarkers, demonstrate the superiority of MagicStain in terms of fidelity, visual quality, and computational efficiency compared to existing methods. Our code and trained models will be released.
Primary Area: applications to computer vision, audio, language, and other modalities
Submission Number: 6476
Loading