MagicStain: High-Fidelity Pathology Image Virtual Staining via Guided Single-Step Diffusion

Virtual staining leverages computational methods to generate different stain styles from a pathology image of a chemically stained tissue section, offering 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 virtual staining remains a significant challenge. While diffusion approaches typically produce more photorealistic images than GAN counterparts through multi-step sampling, this comes at the cost of high computational overhead and inference latency. This paper proposes MagicStain, a novel single-step diffusion model tailored for generating high-resolution virtual stains. Specifically, we adapt a single-step diffusion model to enable efficient virtual staining. By introducing pathology priors from a pathology foundation model and integrating pathology- and structure-consistency losses on both the original images and the hematoxylin channel, MagicStain achieves high-fidelity, 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 at low training cost. Extensive experiments on eight virtual staining datasets, each involving translation between different staining dyes/biomarkers, demonstrate the superiority of MagicStain in fidelity, visual quality, and computational efficiency compared to existing methods. Our code and trained models will be publicly available.