Skip to content

Commit

Permalink
docs: update README.md
Browse files Browse the repository at this point in the history
  • Loading branch information
junhsss committed Mar 27, 2023
1 parent 9257f53 commit cf40702
Showing 1 changed file with 33 additions and 19 deletions.
52 changes: 33 additions & 19 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,29 +10,36 @@

**Consistency Models** are a new family of generative models that achieve high sample quality without adversarial training. They support _fast one-step generation_ by design, while still allowing for few-step sampling to trade compute for sample quality.

<br />

## Installation

```sh
$ pip install consistency
```

### Note

If you just want to try things out, just do:
You **don't need to install** `consistency` for just trying things out:

```python
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline("consistency/cifar10-32-demo", custom_pipeline="consistency/pipeline")
pipeline = DiffusionPipeline.from_pretrained(
"consistency/cifar10-32-demo",
custom_pipeline="consistency/pipeline",
)

pipeline().images[0] # Super Fast Generation! 🀯
```

## Installation

```sh
$ pip install consistency
```
<br />

## Quickstart

### Basic

Just wrap your favorite _U-Net_ with `Consistency`. 😊
Just wrap your favorite _U-Net_ with `Consistency`.

```python
import torch
Expand All @@ -53,7 +60,7 @@ samples = consistency.sample(16, steps=5, use_ema=True)

`Consistency` is self-contained with the training logic and all necessary schedules.

You can train `Consistency` with **PyTorch Lightning**'s `Trainer` πŸš€
You can train it with **PyTorch Lightning**'s `Trainer` πŸš€

```python
from pytorch_lightning import Trainer
Expand All @@ -64,10 +71,6 @@ trainer.fit(consistency, some_dataloader)

<br />

A complete example can be found in [this **script**](https://github.com/junhsss/consistency-models/blob/main/examples/train.py) or in [this **colab notebook**](https://colab.research.google.com/github/junhsss/consistency-models/blob/main/examples/consistency_models.ipynb).

Checkout [this **Wandb workspace**](https://wandb.ai/junhsss/consistency?workspace=user-junhsss) for some experiment results.

### Push to HF Hub

Just provide your `model_id` and `token`!
Expand All @@ -82,16 +85,27 @@ consistency = Consistency(
)
```

You can safely uninstall `consistency` afterwards. Good luck! 🀞:
You can safely drop `consistency` afterwards. Good luck! 🀞

```python
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline("your_model_id", custom_pipeline="consistency/pipeline")
pipeline = DiffusionPipeline.from_pretrained(
"your_model_id",
custom_pipeline="consistency/pipeline",
)

pipeline().images[0]
```

<br />

A complete example can be found in [this **script**](https://github.com/junhsss/consistency-models/blob/main/examples/train.py) or in [this **colab notebook**](https://colab.research.google.com/github/junhsss/consistency-models/blob/main/examples/consistency_models.ipynb).

Checkout [this **Wandb workspace**](https://wandb.ai/junhsss/consistency?workspace=user-junhsss) for some experiment results.

<br />

## Available Models

| model_id | sample_size |
Expand All @@ -100,10 +114,14 @@ pipeline().images[0]

If you've trained some checkpoints. **Share with us! πŸ€—**

<br />

## Documentation

In progress... πŸ› 

<br />

## Reference

```bibtex
Expand All @@ -118,7 +136,3 @@ In progress... πŸ› 
copyright = {arXiv.org perpetual, non-exclusive license}
}
```

## Todo

- [ ] `diffusers` integration. (`ConsistencyPipeline` + `ConsistencyScheduler`)

0 comments on commit cf40702

Please sign in to comment.