TorchDiff 2.5.0__tar.gz → 2.6.0__tar.gz

{torchdiff-2.5.0/TorchDiff.egg-info → torchdiff-2.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TorchDiff
-Version: 2.5.0
+Version: 2.6.0
 Summary: A PyTorch-based library for diffusion models
 Home-page: https://github.com/LoqmanSamani/TorchDiff
 Author: Loghman Samani
@@ -61,7 +61,7 @@ Dynamic: summary
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-red?style=plastic)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-white?style=plastic&logo=pytorch&logoColor=red)](https://pytorch.org/)
-[![Version](https://img.shields.io/badge/version-2.5.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
+[![Version](https://img.shields.io/badge/version-2.6.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
 [![Python](https://img.shields.io/badge/python-3.10%2B-blue?style=plastic&logo=python&logoColor=white)](https://www.python.org/)
 [![Downloads](https://pepy.tech/badge/torchdiff)](https://pepy.tech/project/torchdiff)
 [![Stars](https://img.shields.io/github/stars/LoqmanSamani/TorchDiff?style=plastic&color=yellow)](https://github.com/LoqmanSamani/TorchDiff)
@@ -76,7 +76,7 @@ Dynamic: summary
 
 **TorchDiff** is a PyTorch library for diffusion models, implementing foundational architectures from recent research. The library provides modular components for building, training, and sampling from diffusion-based generative models.
 
-Version 2.5.0 includes five major model families grounded in the diffusion modeling literature. **DDPM** (Ho et al., 2020) and **DDIM** (Song et al., 2021a) establish the core discrete-time framework. **SDE-based diffusion** (Song et al., 2021b) extends this to continuous stochastic processes with variance-exploding and variance-preserving formulations. **LDM** (Rombach et al., 2022) moves diffusion into learned latent spaces via variational autoencoders. **UnCLIP** (Ramesh et al., 2022) combines CLIP embeddings with hierarchical generation for text-to-image synthesis.
+Version 2.6.0 includes five major model families grounded in the diffusion modeling literature. **DDPM** (Ho et al., 2020) and **DDIM** (Song et al., 2021a) establish the core discrete-time framework. **SDE-based diffusion** (Song et al., 2021b) extends this to continuous stochastic processes with variance-exploding and variance-preserving formulations. **LDM** (Rombach et al., 2022) moves diffusion into learned latent spaces via variational autoencoders. **UnCLIP** (Ramesh et al., 2022) combines CLIP embeddings with hierarchical generation for text-to-image synthesis.
 
 <div align="center">
   <img src="https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/imgs/mount.png?raw=true" alt="Diffusion Model Process" width="1000"/>
@@ -93,17 +93,9 @@ We also provide evaluation utilities including standard metrics (MSE, PSNR, SSIM
 
 ---
 
-## What's New in v2.5.0
 
-- **UnCLIP improvements**: Fixed CLIPContextProjection output dimension handling, corrected sampling loop index arithmetic, resolved NaN loss in upsampler/prior training via bfloat16 autocast, and fixed CLIPEmbeddingProjection reconstruction loss bug.
-- **Expanded test coverage**: Added test suites for LDM (AutoencoderLDM), UnCLIP (Scheduler, Forward/Reverse, Projections, TransformerPrior), and Utils (DiffusionNetwork, loss functions, Metrics).
-- **API completeness**: `TrainUnCLIPPrior` now properly exported; removed duplicate `SampleUnCLIP` import.
-- **Documentation**: Aligned all RST titles, added `torchmetrics` to mock imports for ReadTheDocs builds.
-- **Build fixes**: Corrected ReadTheDocs URL in setup.py, removed trailing commas from requirements.txt, unified README for both GitHub and PyPI.
 
----
-
-## Installation
+### Installation
 
 Install the stable release from PyPI.
 
@@ -213,7 +205,8 @@ DDPM (Ho et al., 2020) frames generation as learning to reverse a Markov chain t
 
 The implementation supports both unconditional generation and conditional variants where generation is guided by auxiliary information like class labels or text embeddings.
 
-**Paper:** [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)  
+**Paper:** [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
+
 **Example:** [DDPM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ddpm/ddpm.ipynb)
 
 ---
@@ -224,7 +217,8 @@ DDIM (Song et al., 2021a) reformulates the generative process as a non-Markovian
 
 Like DDPM, both conditional and unconditional generation modes are supported.
 
-**Paper:** [Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)  
+**Paper:** [Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)
+
 **Example:** [DDIM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ddim/ddim.ipynb)
 
 ---
@@ -235,7 +229,8 @@ The SDE framework (Song et al., 2021b) generalizes diffusion models as continuou
 
 We implement variance-exploding (VE), variance-preserving (VP), and sub-VP formulations. The reverse process can be simulated using either stochastic differential equations or their deterministic probability flow ODE counterparts. This unifies score matching with denoising diffusion and enables more flexible sampling strategies.
 
-**Paper:** [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456)  
+**Paper:** [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456)
+
 **Example:** [SDE Notebooks](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/sde/)
 
 ---
@@ -246,7 +241,8 @@ LDM (Rombach et al., 2022) addresses the computational cost of pixel-space diffu
 
 Any of the diffusion backends (DDPM, DDIM, SDE) can operate in this latent space. The architecture enables high-resolution synthesis that would be impractical in pixel space.
 
-**Paper:** [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)  
+**Paper:** [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)
+
 **Example:** [LDM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ldm/ldm.ipynb)
 
 ---
@@ -259,7 +255,8 @@ This hierarchical approach leverages CLIP's multimodal embedding space where tex
 
 Given the complexity, UnCLIP training requires more extensive setup than other models in this library.
 
-**Paper:** [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://arxiv.org/abs/2204.06125)  
+**Paper:** [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://arxiv.org/abs/2204.06125)
+
 **Example:** [UnCLIP Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/unclip/unclip.ipynb)
 
 ---
@@ -268,15 +265,16 @@ Given the complexity, UnCLIP training requires more extensive setup than other m
 
 TorchDiff breaks each model into reusable components:
 
-| Component | Description |
-|-----------|-------------|
-| **Forward Diffusion** | Adds noise to data following model-specific schedules |
-| **Reverse Diffusion** | Removes noise to recover data via learned denoising |
-| **Scheduler** | Controls variance/noise schedules across timesteps |
-| **Training** | Complete training pipelines with mixed precision, gradient accumulation |
-| **Sampling** | Efficient inference and image generation routines |
+| Component                   | Description                                                             |
+| --------------------------- | ----------------------------------------------------------------------- |
+| **Forward Diffusion** | Adds noise to data following model-specific schedules                   |
+| **Reverse Diffusion** | Removes noise to recover data via learned denoising                     |
+| **Scheduler**         | Controls variance/noise schedules across timesteps                      |
+| **Training**          | Complete training pipelines with mixed precision, gradient accumulation |
+| **Sampling**          | Efficient inference and image generation routines                       |
 
 Additional utilities:
+
 - **DiffusionNetwork**: U-Net architecture with attention and time embeddings
 - **TextEncoder**: Transformer-based encoder for conditional generation
 - **Metrics**: Evaluation suite (MSE, PSNR, SSIM, FID, LPIPS)
@@ -297,13 +295,13 @@ Documentation and additional materials are available online.
 
 We are actively developing TorchDiff with several improvements planned for future releases.
 
-**Model Extensions**  
+**Model Extensions**
 New diffusion variants and training algorithms from recent literature will be added as they become established. We are particularly interested in methods that improve sample efficiency or generation quality.
 
-**Performance Optimization**  
+**Performance Optimization**
 Sampling speed and memory efficiency remain active areas of research. We plan to integrate faster sampling methods and more efficient architectures as they emerge.
 
-**Experimental Utilities**  
+**Experimental Utilities**
 Additional tools for hyperparameter tuning, ablation studies, and model comparison will make experimentation more straightforward.
 
 ---

{torchdiff-2.5.0 → torchdiff-2.6.0}/README.md

@@ -8,7 +8,7 @@
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-red?style=plastic)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-white?style=plastic&logo=pytorch&logoColor=red)](https://pytorch.org/)
-[![Version](https://img.shields.io/badge/version-2.5.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
+[![Version](https://img.shields.io/badge/version-2.6.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
 [![Python](https://img.shields.io/badge/python-3.10%2B-blue?style=plastic&logo=python&logoColor=white)](https://www.python.org/)
 [![Downloads](https://pepy.tech/badge/torchdiff)](https://pepy.tech/project/torchdiff)
 [![Stars](https://img.shields.io/github/stars/LoqmanSamani/TorchDiff?style=plastic&color=yellow)](https://github.com/LoqmanSamani/TorchDiff)
@@ -23,7 +23,7 @@
 
 **TorchDiff** is a PyTorch library for diffusion models, implementing foundational architectures from recent research. The library provides modular components for building, training, and sampling from diffusion-based generative models.
 
-Version 2.5.0 includes five major model families grounded in the diffusion modeling literature. **DDPM** (Ho et al., 2020) and **DDIM** (Song et al., 2021a) establish the core discrete-time framework. **SDE-based diffusion** (Song et al., 2021b) extends this to continuous stochastic processes with variance-exploding and variance-preserving formulations. **LDM** (Rombach et al., 2022) moves diffusion into learned latent spaces via variational autoencoders. **UnCLIP** (Ramesh et al., 2022) combines CLIP embeddings with hierarchical generation for text-to-image synthesis.
+Version 2.6.0 includes five major model families grounded in the diffusion modeling literature. **DDPM** (Ho et al., 2020) and **DDIM** (Song et al., 2021a) establish the core discrete-time framework. **SDE-based diffusion** (Song et al., 2021b) extends this to continuous stochastic processes with variance-exploding and variance-preserving formulations. **LDM** (Rombach et al., 2022) moves diffusion into learned latent spaces via variational autoencoders. **UnCLIP** (Ramesh et al., 2022) combines CLIP embeddings with hierarchical generation for text-to-image synthesis.
 
 <div align="center">
   <img src="https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/imgs/mount.png?raw=true" alt="Diffusion Model Process" width="1000"/>
@@ -40,17 +40,9 @@ We also provide evaluation utilities including standard metrics (MSE, PSNR, SSIM
 
 ---
 
-## What's New in v2.5.0
 
-- **UnCLIP improvements**: Fixed CLIPContextProjection output dimension handling, corrected sampling loop index arithmetic, resolved NaN loss in upsampler/prior training via bfloat16 autocast, and fixed CLIPEmbeddingProjection reconstruction loss bug.
-- **Expanded test coverage**: Added test suites for LDM (AutoencoderLDM), UnCLIP (Scheduler, Forward/Reverse, Projections, TransformerPrior), and Utils (DiffusionNetwork, loss functions, Metrics).
-- **API completeness**: `TrainUnCLIPPrior` now properly exported; removed duplicate `SampleUnCLIP` import.
-- **Documentation**: Aligned all RST titles, added `torchmetrics` to mock imports for ReadTheDocs builds.
-- **Build fixes**: Corrected ReadTheDocs URL in setup.py, removed trailing commas from requirements.txt, unified README for both GitHub and PyPI.
 
----
-
-## Installation
+### Installation
 
 Install the stable release from PyPI.
 
@@ -160,7 +152,8 @@ DDPM (Ho et al., 2020) frames generation as learning to reverse a Markov chain t
 
 The implementation supports both unconditional generation and conditional variants where generation is guided by auxiliary information like class labels or text embeddings.
 
-**Paper:** [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)  
+**Paper:** [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
+
 **Example:** [DDPM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ddpm/ddpm.ipynb)
 
 ---
@@ -171,7 +164,8 @@ DDIM (Song et al., 2021a) reformulates the generative process as a non-Markovian
 
 Like DDPM, both conditional and unconditional generation modes are supported.
 
-**Paper:** [Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)  
+**Paper:** [Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)
+
 **Example:** [DDIM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ddim/ddim.ipynb)
 
 ---
@@ -182,7 +176,8 @@ The SDE framework (Song et al., 2021b) generalizes diffusion models as continuou
 
 We implement variance-exploding (VE), variance-preserving (VP), and sub-VP formulations. The reverse process can be simulated using either stochastic differential equations or their deterministic probability flow ODE counterparts. This unifies score matching with denoising diffusion and enables more flexible sampling strategies.
 
-**Paper:** [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456)  
+**Paper:** [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456)
+
 **Example:** [SDE Notebooks](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/sde/)
 
 ---
@@ -193,7 +188,8 @@ LDM (Rombach et al., 2022) addresses the computational cost of pixel-space diffu
 
 Any of the diffusion backends (DDPM, DDIM, SDE) can operate in this latent space. The architecture enables high-resolution synthesis that would be impractical in pixel space.
 
-**Paper:** [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)  
+**Paper:** [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)
+
 **Example:** [LDM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ldm/ldm.ipynb)
 
 ---
@@ -206,7 +202,8 @@ This hierarchical approach leverages CLIP's multimodal embedding space where tex
 
 Given the complexity, UnCLIP training requires more extensive setup than other models in this library.
 
-**Paper:** [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://arxiv.org/abs/2204.06125)  
+**Paper:** [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://arxiv.org/abs/2204.06125)
+
 **Example:** [UnCLIP Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/unclip/unclip.ipynb)
 
 ---
@@ -215,15 +212,16 @@ Given the complexity, UnCLIP training requires more extensive setup than other m
 
 TorchDiff breaks each model into reusable components:
 
-| Component | Description |
-|-----------|-------------|
-| **Forward Diffusion** | Adds noise to data following model-specific schedules |
-| **Reverse Diffusion** | Removes noise to recover data via learned denoising |
-| **Scheduler** | Controls variance/noise schedules across timesteps |
-| **Training** | Complete training pipelines with mixed precision, gradient accumulation |
-| **Sampling** | Efficient inference and image generation routines |
+| Component                   | Description                                                             |
+| --------------------------- | ----------------------------------------------------------------------- |
+| **Forward Diffusion** | Adds noise to data following model-specific schedules                   |
+| **Reverse Diffusion** | Removes noise to recover data via learned denoising                     |
+| **Scheduler**         | Controls variance/noise schedules across timesteps                      |
+| **Training**          | Complete training pipelines with mixed precision, gradient accumulation |
+| **Sampling**          | Efficient inference and image generation routines                       |
 
 Additional utilities:
+
 - **DiffusionNetwork**: U-Net architecture with attention and time embeddings
 - **TextEncoder**: Transformer-based encoder for conditional generation
 - **Metrics**: Evaluation suite (MSE, PSNR, SSIM, FID, LPIPS)
@@ -244,13 +242,13 @@ Documentation and additional materials are available online.
 
 We are actively developing TorchDiff with several improvements planned for future releases.
 
-**Model Extensions**  
+**Model Extensions**
 New diffusion variants and training algorithms from recent literature will be added as they become established. We are particularly interested in methods that improve sample efficiency or generation quality.
 
-**Performance Optimization**  
+**Performance Optimization**
 Sampling speed and memory efficiency remain active areas of research. We plan to integrate faster sampling methods and more efficient architectures as they emerge.
 
-**Experimental Utilities**  
+**Experimental Utilities**
 Additional tools for hyperparameter tuning, ablation studies, and model comparison will make experimentation more straightforward.
 
 ---

{torchdiff-2.5.0 → torchdiff-2.6.0/TorchDiff.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: TorchDiff
-Version: 2.5.0
+Version: 2.6.0
 Summary: A PyTorch-based library for diffusion models
 Home-page: https://github.com/LoqmanSamani/TorchDiff
 Author: Loghman Samani
@@ -61,7 +61,7 @@ Dynamic: summary
 
 [![License: MIT](https://img.shields.io/badge/license-MIT-red?style=plastic)](https://opensource.org/licenses/MIT)
 [![PyTorch](https://img.shields.io/badge/PyTorch-white?style=plastic&logo=pytorch&logoColor=red)](https://pytorch.org/)
-[![Version](https://img.shields.io/badge/version-2.5.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
+[![Version](https://img.shields.io/badge/version-2.6.0-blue?style=plastic)](https://pypi.org/project/torchdiff/)
 [![Python](https://img.shields.io/badge/python-3.10%2B-blue?style=plastic&logo=python&logoColor=white)](https://www.python.org/)
 [![Downloads](https://pepy.tech/badge/torchdiff)](https://pepy.tech/project/torchdiff)
 [![Stars](https://img.shields.io/github/stars/LoqmanSamani/TorchDiff?style=plastic&color=yellow)](https://github.com/LoqmanSamani/TorchDiff)
@@ -76,7 +76,7 @@ Dynamic: summary
 
 **TorchDiff** is a PyTorch library for diffusion models, implementing foundational architectures from recent research. The library provides modular components for building, training, and sampling from diffusion-based generative models.
 
-Version 2.5.0 includes five major model families grounded in the diffusion modeling literature. **DDPM** (Ho et al., 2020) and **DDIM** (Song et al., 2021a) establish the core discrete-time framework. **SDE-based diffusion** (Song et al., 2021b) extends this to continuous stochastic processes with variance-exploding and variance-preserving formulations. **LDM** (Rombach et al., 2022) moves diffusion into learned latent spaces via variational autoencoders. **UnCLIP** (Ramesh et al., 2022) combines CLIP embeddings with hierarchical generation for text-to-image synthesis.
+Version 2.6.0 includes five major model families grounded in the diffusion modeling literature. **DDPM** (Ho et al., 2020) and **DDIM** (Song et al., 2021a) establish the core discrete-time framework. **SDE-based diffusion** (Song et al., 2021b) extends this to continuous stochastic processes with variance-exploding and variance-preserving formulations. **LDM** (Rombach et al., 2022) moves diffusion into learned latent spaces via variational autoencoders. **UnCLIP** (Ramesh et al., 2022) combines CLIP embeddings with hierarchical generation for text-to-image synthesis.
 
 <div align="center">
   <img src="https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/imgs/mount.png?raw=true" alt="Diffusion Model Process" width="1000"/>
@@ -93,17 +93,9 @@ We also provide evaluation utilities including standard metrics (MSE, PSNR, SSIM
 
 ---
 
-## What's New in v2.5.0
 
-- **UnCLIP improvements**: Fixed CLIPContextProjection output dimension handling, corrected sampling loop index arithmetic, resolved NaN loss in upsampler/prior training via bfloat16 autocast, and fixed CLIPEmbeddingProjection reconstruction loss bug.
-- **Expanded test coverage**: Added test suites for LDM (AutoencoderLDM), UnCLIP (Scheduler, Forward/Reverse, Projections, TransformerPrior), and Utils (DiffusionNetwork, loss functions, Metrics).
-- **API completeness**: `TrainUnCLIPPrior` now properly exported; removed duplicate `SampleUnCLIP` import.
-- **Documentation**: Aligned all RST titles, added `torchmetrics` to mock imports for ReadTheDocs builds.
-- **Build fixes**: Corrected ReadTheDocs URL in setup.py, removed trailing commas from requirements.txt, unified README for both GitHub and PyPI.
 
----
-
-## Installation
+### Installation
 
 Install the stable release from PyPI.
 
@@ -213,7 +205,8 @@ DDPM (Ho et al., 2020) frames generation as learning to reverse a Markov chain t
 
 The implementation supports both unconditional generation and conditional variants where generation is guided by auxiliary information like class labels or text embeddings.
 
-**Paper:** [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)  
+**Paper:** [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
+
 **Example:** [DDPM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ddpm/ddpm.ipynb)
 
 ---
@@ -224,7 +217,8 @@ DDIM (Song et al., 2021a) reformulates the generative process as a non-Markovian
 
 Like DDPM, both conditional and unconditional generation modes are supported.
 
-**Paper:** [Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)  
+**Paper:** [Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)
+
 **Example:** [DDIM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ddim/ddim.ipynb)
 
 ---
@@ -235,7 +229,8 @@ The SDE framework (Song et al., 2021b) generalizes diffusion models as continuou
 
 We implement variance-exploding (VE), variance-preserving (VP), and sub-VP formulations. The reverse process can be simulated using either stochastic differential equations or their deterministic probability flow ODE counterparts. This unifies score matching with denoising diffusion and enables more flexible sampling strategies.
 
-**Paper:** [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456)  
+**Paper:** [Score-Based Generative Modeling through Stochastic Differential Equations](https://arxiv.org/abs/2011.13456)
+
 **Example:** [SDE Notebooks](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/sde/)
 
 ---
@@ -246,7 +241,8 @@ LDM (Rombach et al., 2022) addresses the computational cost of pixel-space diffu
 
 Any of the diffusion backends (DDPM, DDIM, SDE) can operate in this latent space. The architecture enables high-resolution synthesis that would be impractical in pixel space.
 
-**Paper:** [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)  
+**Paper:** [High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752)
+
 **Example:** [LDM Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/ldm/ldm.ipynb)
 
 ---
@@ -259,7 +255,8 @@ This hierarchical approach leverages CLIP's multimodal embedding space where tex
 
 Given the complexity, UnCLIP training requires more extensive setup than other models in this library.
 
-**Paper:** [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://arxiv.org/abs/2204.06125)  
+**Paper:** [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://arxiv.org/abs/2204.06125)
+
 **Example:** [UnCLIP Notebook](https://github.com/LoqmanSamani/TorchDiff/blob/systembiology/examples/unclip/unclip.ipynb)
 
 ---
@@ -268,15 +265,16 @@ Given the complexity, UnCLIP training requires more extensive setup than other m
 
 TorchDiff breaks each model into reusable components:
 
-| Component | Description |
-|-----------|-------------|
-| **Forward Diffusion** | Adds noise to data following model-specific schedules |
-| **Reverse Diffusion** | Removes noise to recover data via learned denoising |
-| **Scheduler** | Controls variance/noise schedules across timesteps |
-| **Training** | Complete training pipelines with mixed precision, gradient accumulation |
-| **Sampling** | Efficient inference and image generation routines |
+| Component                   | Description                                                             |
+| --------------------------- | ----------------------------------------------------------------------- |
+| **Forward Diffusion** | Adds noise to data following model-specific schedules                   |
+| **Reverse Diffusion** | Removes noise to recover data via learned denoising                     |
+| **Scheduler**         | Controls variance/noise schedules across timesteps                      |
+| **Training**          | Complete training pipelines with mixed precision, gradient accumulation |
+| **Sampling**          | Efficient inference and image generation routines                       |
 
 Additional utilities:
+
 - **DiffusionNetwork**: U-Net architecture with attention and time embeddings
 - **TextEncoder**: Transformer-based encoder for conditional generation
 - **Metrics**: Evaluation suite (MSE, PSNR, SSIM, FID, LPIPS)
@@ -297,13 +295,13 @@ Documentation and additional materials are available online.
 
 We are actively developing TorchDiff with several improvements planned for future releases.
 
-**Model Extensions**  
+**Model Extensions**
 New diffusion variants and training algorithms from recent literature will be added as they become established. We are particularly interested in methods that improve sample efficiency or generation quality.
 
-**Performance Optimization**  
+**Performance Optimization**
 Sampling speed and memory efficiency remain active areas of research. We plan to integrate faster sampling methods and more efficient architectures as they emerge.
 
-**Experimental Utilities**  
+**Experimental Utilities**
 Additional tools for hyperparameter tuning, ablation studies, and model comparison will make experimentation more straightforward.
 
 ---

{torchdiff-2.5.0 → torchdiff-2.6.0}/TorchDiff.egg-info/SOURCES.txt

@@ -41,7 +41,17 @@ torchdiff/sde.py
 torchdiff/unclip.py
 torchdiff/utils.py
 torchdiff/tests/__init__.py
+torchdiff/tests/bench_ddim.py
+torchdiff/tests/bench_ddpm.py
+torchdiff/tests/bench_ldm.py
+torchdiff/tests/bench_sde.py
+torchdiff/tests/bench_unclip.py
 torchdiff/tests/test_ddim.py
+torchdiff/tests/test_ddp_ddim.py
+torchdiff/tests/test_ddp_ddpm.py
+torchdiff/tests/test_ddp_ldm.py
+torchdiff/tests/test_ddp_sde.py
+torchdiff/tests/test_ddp_unclip.py
 torchdiff/tests/test_ddpm.py
 torchdiff/tests/test_ldm.py
 torchdiff/tests/test_sde.py

{torchdiff-2.5.0 → torchdiff-2.6.0}/setup.py

@@ -11,7 +11,7 @@ if not long_description:
 
 setup(
     name="TorchDiff",
-    version="2.5.0",
+    version="2.6.0",
     description="A PyTorch-based library for diffusion models",
     long_description=long_description,
     long_description_content_type="text/markdown",

{torchdiff-2.5.0 → torchdiff-2.6.0}/torchdiff/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.5.0"
+__version__ = "2.6.0"
 
 from .ddim import ForwardDDIM, ReverseDDIM, SchedulerDDIM, TrainDDIM, SampleDDIM
 from .ddpm import ForwardDDPM, ReverseDDPM, SchedulerDDPM, TrainDDPM, SampleDDPM

{torchdiff-2.5.0 → torchdiff-2.6.0}/torchdiff/ddim.py

@@ -46,6 +46,14 @@ from typing_extensions import Self
 from .utils import LossAdapter
 import os
 
+__all__ = [
+    "ForwardDDIM",
+    "ReverseDDIM",
+    "SchedulerDDIM",
+    "TrainDDIM",
+    "SampleDDIM",
+]
+
 
 ###==================================================================================================================###
 
@@ -376,8 +384,7 @@ class SchedulerDDIM(nn.Module):
             Reshaped tensor suitable for broadcasting.
         """
         batch_size = t.shape[0]
-        out = t.to(t.device)
-        return out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))
+        return t.reshape(batch_size, *((1,) * (len(x_shape) - 1)))
 
 
 ###==================================================================================================================###
@@ -506,6 +513,7 @@ class TrainDDIM(nn.Module):
             factor=0.5
         )
         self.warmup_lr_scheduler = self.warmup_scheduler(self.optim, warmup_steps)
+        self._device_type = self.device.type if hasattr(self.device, 'type') else ('cuda' if 'cuda' in str(self.device) else 'cpu')
         if tokenizer is None:
             try:
                 self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
@@ -524,15 +532,17 @@ class TrainDDIM(nn.Module):
             raise ValueError("DDP enabled but LOCAL_RANK environment variable not set")
         if "WORLD_SIZE" not in os.environ:
             raise ValueError("DDP enabled but WORLD_SIZE environment variable not set")
-        if not torch.cuda.is_available():
-            raise RuntimeError("DDP requires CUDA but CUDA is not available")
         if not torch.distributed.is_initialized():
-            init_process_group(backend="nccl")
+            backend = "nccl" if torch.cuda.is_available() else "gloo"
+            init_process_group(backend=backend)
         self.ddp_rank = int(os.environ["RANK"])  # global rank across all nodes
         self.ddp_local_rank = int(os.environ["LOCAL_RANK"])  # local rank on current node
         self.ddp_world_size = int(os.environ["WORLD_SIZE"])  # total number of processes
-        self.device = torch.device(f"cuda:{self.ddp_local_rank}")
-        torch.cuda.set_device(self.device)
+        if torch.cuda.is_available():
+            self.device = torch.device(f"cuda:{self.ddp_local_rank}")
+            torch.cuda.set_device(self.device)
+        else:
+            self.device = torch.device("cpu")
         self.master_process = self.ddp_rank == 0
         if self.master_process:
             print(f"DDP initialized with world_size={self.ddp_world_size}")
@@ -641,17 +651,12 @@ class TrainDDIM(nn.Module):
     def _wrap_models_for_ddp(self) -> None:
         """Wrap models with DistributedDataParallel for multi-GPU training."""
         if self.use_ddp:
-            self.diff_net = DDP(
-                self.diff_net,
-                device_ids=[self.ddp_local_rank],
-                find_unused_parameters=True
-            )
+            ddp_kwargs = dict(find_unused_parameters=False)
+            if self._device_type == 'cuda':
+                ddp_kwargs['device_ids'] = [self.ddp_local_rank]
+            self.diff_net = DDP(self.diff_net, **ddp_kwargs)
             if self.cond_net is not None:
-                self.cond_net = DDP(
-                    self.cond_net,
-                    device_ids=[self.ddp_local_rank],
-                    find_unused_parameters=True
-                )
+                self.cond_net = DDP(self.cond_net, **ddp_kwargs)
 
     def forward(self) -> Dict:
         """Trains the DDIM model to predict noise added by the forward diffusion process.
@@ -678,7 +683,10 @@ class TrainDDIM(nn.Module):
                     print(f"Model compilation failed: {e}. Continuing without compilation.")
 
         self._wrap_models_for_ddp()
-        scaler = torch.GradScaler()
+        use_amp = self._device_type == 'cuda'
+        scaler = torch.amp.GradScaler(self._device_type, enabled=use_amp)
+        if use_amp:
+            torch.backends.cudnn.benchmark = True
         wait = 0
         for epoch in range(self.max_epochs):
             pbar = tqdm(self.train_loader, desc=f"Epoch {epoch + 1}/{self.max_epochs}", disable=not self.master_process)
@@ -691,7 +699,7 @@ class TrainDDIM(nn.Module):
                     y_encoded = self._process_conditional_input(y)
                 else:
                     y_encoded = None
-                with torch.autocast(device_type='cuda' if self.device == 'cuda' else 'cpu'):
+                with torch.autocast(device_type=self._device_type, enabled=use_amp):
                     noise = torch.randn_like(x)
                     t = torch.randint(0, self.fwd_ddim.vs.train_steps, (x.shape[0],), device=x.device)
                     xt, target = self.fwd_ddim(x, t, noise)
@@ -706,8 +714,8 @@ class TrainDDIM(nn.Module):
                         torch.nn.utils.clip_grad_norm_(self.cond_net.parameters(), max_norm=1.0)
                     scaler.step(self.optim)
                     scaler.update()
-                    self.optim.zero_grad()
-                    if self.global_step > 0 and self.global_step < self.warmup_steps:
+                    self.optim.zero_grad(set_to_none=True)
+                    if self.global_step < self.warmup_steps:
                         self.warmup_lr_scheduler.step()
                     self.global_step += 1
                 pbar.set_postfix({'Loss': f'{loss.item() * self.grad_acc:.4f}'})
@@ -1042,7 +1050,7 @@ class SampleDDIM(nn.Module):
         if conds is None and self.cond_net is not None:
             raise ValueError("Conditions must be provided for conditional model")
 
-        init_samps = torch.randn(self.batch_size, self.in_channels, self.img_size[0], self.img_size[1]).to(self.device)
+        init_samps = torch.randn(self.batch_size, self.in_channels, self.img_size[0], self.img_size[1], device=self.device)
         self.diff_net.eval()
         if self.cond_net:
             self.cond_net.eval()
@@ -1055,14 +1063,13 @@ class SampleDDIM(nn.Module):
             dynamic_ncols=True,
             leave=True,
         )
-        if self.cond_net is not None and conds is not None:
-            input_ids, attention_masks = self.tokenize(conds)
-            key_padding_mask = (attention_masks == 0)
-            y = self.cond_net(input_ids, key_padding_mask)
-        else:
-            y = None
-
         with torch.no_grad():
+            if self.cond_net is not None and conds is not None:
+                input_ids, attention_masks = self.tokenize(conds)
+                key_padding_mask = (attention_masks == 0)
+                y = self.cond_net(input_ids, key_padding_mask)
+            else:
+                y = None
             xt = init_samps
             for i in iterator:
                 t_current = timesteps[i].item()
@@ -1099,6 +1106,7 @@ class SampleDDIM(nn.Module):
         """
         self.device = device
         self.diff_net.to(device)
+        self.rwd_ddim.to(device)
         if self.cond_net:
             self.cond_net.to(device)
         return super().to(device)