Update for tests (#3)

* [WIP] add files

* rename unet

* rename scheduler

* update

* update

* update

Author: shunk031

.gitignore

@@ -1,6 +1,10 @@
 # Created by https://www.toptal.com/developers/gitignore/api/python
 # Edit at https://www.toptal.com/developers/gitignore?templates=python
 
+*.png
+*.json
+*.safetensors
+
 ### Python ###
 # Byte-compiled / optimized / DLL files
 __pycache__/

README.md

@@ -1,24 +1,104 @@
 # 🤗 Noise Conditional Score Networks
 
-[![CI](https://github.com/py-img-gen/diffusers-ncsn/actions/workflows/ci.yaml/badge.svg)](https://github.com/py-img-gen/diffusers-ncsn/actions/workflows/ci.yaml) [![](https://img.shields.io/badge/Official_code-GitHub-green)](https://github.com/ermongroup/ncsn)
+[![CI](https://github.com/py-img-gen/diffusers-ncsn/actions/workflows/ci.yaml/badge.svg)](https://github.com/py-img-gen/diffusers-ncsn/actions/workflows/ci.yaml) 
+[![](https://img.shields.io/badge/Official_code-GitHub-green)](https://github.com/ermongroup/ncsn)
+[![Model on HF](https://img.shields.io/badge/🤗%20Model%20on%20HF-py--img--gen/ncsn--mnist-D4AA00)](https://huggingface.co/py-img-gen/ncsn-mnist)
 
 [`🤗 diffusers`](https://github.com/huggingface/diffusers) implementation of the paper ["Generative Modeling by Estimating Gradients of the Data Distribution" [Yang+ NeurIPS'19]](https://arxiv.org/abs/1907.05600).
 
-## Installation
+## How to use
+
+### Use without installation
+
+You can load the pretrained pipeline directly from the HF Hub as follows:
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.utils import make_image_grid
+
+# Specify the device to use
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+#
+# Load the pipeline from the Hugging Face Hub
+#
+pipe = DiffusionPipeline.from_pretrained(
+    "py-img-gen/ncsn-mnist", trust_remote_code=True
+)
+pipe = pipe.to(device)
+
+# Generate samples; here, we specify the seed and generate 16 images
+output = pipe(
+    batch_size=16,
+    generator=torch.manual_seed(42),
+)
+
+# Create a grid image from the generated samples
+image = make_image_grid(images=output.images, rows=4, cols=4)
+image.save("output.png")
+```
+
+### Use with installation
+
+First, install the package from this repository:
 
 ```shell
 pip install git+https://github.com/py-img-gen/diffusers-ncsn
 ```
 
+Then, you can use the package as follows:
+
+```python
+import torch
+
+from ncsn.pipeline_ncsn import NCSNPipeline
+
+# Specify the device to use
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+#
+# Load the pipeline from the HF Hub through the NCSNPipeline of this library
+#
+pipe = NCSNPipeline.from_pretrained("py-img-gen/ncsn-mnist", trust_remote_code=True)
+pipe = pipe.to(device)
+
+# Generate samples; here, we specify the seed and generate 16 images
+output = pipe(
+    batch_size=16,
+    generator=torch.manual_seed(42),
+)
+
+# Create a grid image from the generated samples
+image = make_image_grid(images=output.images, rows=4, cols=4)
+image.save("output.png")
+```
+
+## Pretrained models and pipeline
+
+[![Model on HF](https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-sm.svg)](https://huggingface.co/py-img-gen/ncsn-mnist) 
+
 ## Showcase
 
 ### MNIST
 
 Example of generating MNIST character images using the model trained with [`train_mnist.py`](https://github.com/py-img-gen/diffusers-ncsn/blob/main/train_mnist.py).
 
-<div align="center">
+<p align="center">
     <img alt="mnist" src="https://github.com/user-attachments/assets/483b6637-2684-4844-8aa1-12b866d46226" width="50%" />
-</div>
+</p>
+
+# Notes on uploading pipelines to the HF Hub with custom components
+
+While referring to 📝 [Load community pipelines and components - huggingface diffusers](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview#community-components), pay attention to the following points.
+- Change [the `_class_name` attribute](https://huggingface.co/py-img-gen/ncsn-mnist/blob/main/model_index.json#L2) in [`model_index.json`](https://huggingface.co/py-img-gen/ncsn-mnist/blob/main/model_index.json) to `["pipeline_ncsn", "NCSNPipeline"]`.
+- Upload [`pipeline_ncsn.py`](https://github.com/py-img-gen/diffusers-ncsn/blob/main/src/ncsn/pipeline_ncsn.py) to [the root of the pipeline repository](https://huggingface.co/py-img-gen/ncsn-mnist/blob/main/pipeline_ncsn.py).
+- Upload custom components to each subfolder:
+  - Upload [`scheduling_ncsn.py`](https://github.com/py-img-gen/diffusers-ncsn/blob/main/src/ncsn/scheduler/scheduling_ncsn.py) to [the `scheduler` subfolder](https://huggingface.co/py-img-gen/ncsn-mnist/tree/main/scheduler).
+  - Upload [`unet_2d_ncsn.py`](https://github.com/py-img-gen/diffusers-ncsn/blob/main/src/ncsn/unet/unet_2d_ncsn.py) to [the `unet` subfolder](https://huggingface.co/py-img-gen/ncsn-mnist/tree/main/unet).
+- Ensure that the custom components are placed in each subfolder because they are referenced by relative paths from `pipeline_ncsn.py`.
+  - Based on this, the code in this library is also placed in the same directory structure as the HF Hub.
+  - For example, `pipeline_ncsn.py` imports `unet_2d_ncsn.py` as `from .unet.unet_2d_ncsn import UNet2DModelForNCSN` because it is placed in the `unet` subfolder.
 
 ## Acknowledgements
 

src/ncsn/pipeline_ncsn.py

@@ -5,11 +5,11 @@
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
 from einops import rearrange
 
-from .scheduling_ncsn import (
-    AnnealedLangevinDynamicOutput,
-    AnnealedLangevinDynamicScheduler,
+from .scheduler.scheduling_ncsn import (
+    AnnealedLangevinDynamicsOutput,
+    AnnealedLangevinDynamicsScheduler,
 )
-from .unet_2d_ncsn import UNet2DModelForNCSN
+from .unet.unet_2d_ncsn import UNet2DModelForNCSN
 
 
 def normalize_images(image: torch.Tensor) -> torch.Tensor:
@@ -42,17 +42,17 @@ class NCSNPipeline(DiffusionPipeline):
     Parameters:
         unet ([`UNet2DModelForNCSN`]):
             A `UNet2DModelForNCSN` to estimate the score of the image.
-        scheduler ([`AnnealedLangevinDynamicScheduler`]):
-            A `AnnealedLangevinDynamicScheduler` to be used in combination with `unet` to estimate the score of the image.
+        scheduler ([`AnnealedLangevinDynamicsScheduler`]):
+            A `AnnealedLangevinDynamicsScheduler` to be used in combination with `unet` to estimate the score of the image.
     """
 
     unet: UNet2DModelForNCSN
-    scheduler: AnnealedLangevinDynamicScheduler
+    scheduler: AnnealedLangevinDynamicsScheduler
 
     _callback_tensor_inputs: List[str] = ["samples"]
 
     def __init__(
-        self, unet: UNet2DModelForNCSN, scheduler: AnnealedLangevinDynamicScheduler
+        self, unet: UNet2DModelForNCSN, scheduler: AnnealedLangevinDynamicsScheduler
     ) -> None:
         super().__init__()
         self.register_modules(unet=unet, scheduler=scheduler)
@@ -151,7 +151,7 @@ def __call__(
                 )
                 samples = (
                     output.prev_sample
-                    if isinstance(output, AnnealedLangevinDynamicOutput)
+                    if isinstance(output, AnnealedLangevinDynamicsOutput)
                     else output[0]
                 )
 

src/ncsn/scheduler/__init__.py

@@ -0,0 +1,9 @@
+from .scheduling_ncsn import (
+    AnnealedLangevinDynamicsOutput,
+    AnnealedLangevinDynamicsScheduler,
+)
+
+__all__ = [
+    "AnnealedLangevinDynamicsOutput",
+    "AnnealedLangevinDynamicsScheduler",
+]

src/ncsn/scheduling_ncsn.py renamed to src/ncsn/scheduler/scheduling_ncsn.py

@@ -15,11 +15,13 @@
 
 
 @dataclass
-class AnnealedLangevinDynamicOutput(SchedulerOutput):
-    """Annealed Langevin Dynamic output class."""
+class AnnealedLangevinDynamicsOutput(SchedulerOutput):
+    """Annealed Langevin Dynamics output class."""
 
 
-class AnnealedLangevinDynamicScheduler(SchedulerMixin, ConfigMixin):  # type: ignore
+class AnnealedLangevinDynamicsScheduler(SchedulerMixin, ConfigMixin):  # type: ignore
+    """Annealed Langevin Dynamics scheduler for Noise Conditional Score Network (NCSN)."""
+
     order = 1
 
     @register_to_config
@@ -106,13 +108,13 @@ def step(
         samples: torch.Tensor,
         return_dict: bool = True,
         **kwargs,
-    ) -> Union[AnnealedLangevinDynamicOutput, Tuple]:
+    ) -> Union[AnnealedLangevinDynamicsOutput, Tuple]:
         z = torch.randn_like(samples)
         step_size = self.step_size[timestep]
         samples = samples + 0.5 * step_size * model_output + torch.sqrt(step_size) * z
 
         if return_dict:
-            return AnnealedLangevinDynamicOutput(prev_sample=samples)
+            return AnnealedLangevinDynamicsOutput(prev_sample=samples)
         else:
             return (samples,)
 

src/ncsn/unet/__init__.py

@@ -0,0 +1,5 @@
+from .unet_2d_ncsn import UNet2DModelForNCSN
+
+__all__ = [
+    "UNet2DModelForNCSN",
+]

src/ncsn/unet_2d_ncsn.py renamed to src/ncsn/unet/unet_2d_ncsn.py

@@ -0,0 +1,31 @@
+import pathlib
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def device() -> torch.device:
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+@pytest.fixture
+def seed() -> int:
+    return 19950815
+
+
+@pytest.fixture
+def root_dir() -> pathlib.Path:
+    return pathlib.Path(__file__).parents[1]
+
+
+@pytest.fixture
+def project_dir(root_dir: pathlib.Path) -> pathlib.Path:
+    dirpath = root_dir / "outputs"
+    dirpath.mkdir(parents=True, exist_ok=True)
+    return dirpath
+
+
+@pytest.fixture
+def lib_dir(root_dir: pathlib.Path) -> pathlib.Path:
+    return root_dir / "src" / "ncsn"