[Feature] register models and scheduelrs from diffusers

[zengyh1900]

Thanks for your contribution and we appreciate it a lot. The following instructions would make your pull request more healthy and more easily get feedback. If you do not understand some items, don't worry, just make the pull request and seek help from maintainers.

Motivation

1. register models from diffusers into MODELS
2. register schedulers from diffusers into DIFFUSION_SCHEDULERS
3. the registered name is set ad Diffusers+ the class name.

Modification

After register, the models and schedulers from diffusers can be used as below.

Schedulers

1. read the original image

import PIL.Image
import numpy as np

def display_sample(sample, i):
    image_processed = sample.cpu().permute(0, 2, 3, 1)
    image_processed = (image_processed + 1.0) * 127.5
    image_processed = image_processed.numpy().astype(np.uint8)

    image_pil = PIL.Image.fromarray(image_processed[0])
    display(f"Image at step {i}")
    display(image_pil)

from PIL import Image 
img = Image.open('walrus.png').convert('RGB')

2. Get an instance of DDPMScheduler through mmedit.registry

from mmedit.registry import DIFFUSION_SCHEDULERS

scheduler_cfg = dict(
    type = 'DDPMScheduler', 
    num_train_timesteps = 1000,
)

my_scheduler = DIFFUSION_SCHEDULERS.build(scheduler_cfg)

3. convert the image to torch.tensor and sample noise

import numpy as np 

image_latents = np.array(img).astype(np.float32) / 255.0
image_latents = image_latents[None].transpose(0, 3, 1, 2)
image_latents = torch.from_numpy(image_latents)
image_latents = (2.0 * image_latents - 1.0).to('cuda')

# sample a noise 
noise = torch.randn(image_latents.shape).to(image_latents.device)

# we use timestep as 800 for example 
timesteps = torch.LongTensor([800]).to(image_latents.device)

4. using my_scheduler to add noise.

noisy_latents = my_scheduler.add_noise(image_latents, noise, timesteps)
display_sample(noisy_latents, 800)

5. using my_scheduler to remove noise.

out = my_scheduler.step(noise, 800, noisy_latents).pred_original_sample
display_sample(out, 0)

Models

1. get an instance of denoising diffusion model

from mmedit.registry import MODELS 

model_cfg = dict(
    type = 'UNet2DModel',
)

unet = MODELS.build(model_cfg)

2. load pre-trained weights

# here, we load weights from huggingface
repo_id = "google/ddpm-church-256"
unet = unet.from_pretrained(repo_id).to('cuda')

3. show the denoising process step-by-step.

import torch
import tqdm

torch.manual_seed(0)

sample = torch.randn(
    1, unet.config.in_channels, unet.config.sample_size, unet.config.sample_size
)

sample = sample.to('cuda')

for i, t in enumerate(tqdm.tqdm(my_scheduler.timesteps)):
  # 1. predict noise residual
  with torch.no_grad():
      noisy_residual = unet(sample, t).sample

  # 2. compute less noisy image and set x_t -> x_t-1
  sample = my_scheduler.step(noisy_residual, t, sample).prev_sample

  # 3. optionally look at image
  if (i + 1) % 200 == 0:
      display_sample(sample, i + 1)

Who can help? @ them here!

BC-breaking (Optional)

Does the modification introduce changes that break the backward-compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.

Use cases (Optional)

If this PR introduces a new feature, it is better to list some use cases here, and update the documentation.

Checklist

Before PR:

• I have read and followed the workflow indicated in the CONTRIBUTING.md to create this PR.
• Pre-commit or linting tools indicated in CONTRIBUTING.md are used to fix the potential lint issues.
• Bug fixes are covered by unit tests, the case that causes the bug should be added in the unit tests.
• New functionalities are covered by complete unit tests. If not, please add more unit test to ensure the correctness.
• The documentation has been modified accordingly, including docstring or example tutorials.

After PR:

• If the modification has potential influence on downstream or other related projects, this PR should be tested with some of those projects.
• CLA has been signed and all committers have signed the CLA in this PR.

[codecov]

Codecov Report

Patch coverage: 84.09% and project coverage change: -0.01 ⚠️

Comparison is base (4a9dfa4) 88.15% compared to head (7ed5724) 88.14%.

❗ Current head 7ed5724 differs from pull request most recent head 18a4939. Consider uploading reports for the commit 18a4939 to get more accurate results

Additional details and impacted files

@@             Coverage Diff             @@
##           dev-1.x    #1692      +/-   ##
===========================================
- Coverage    88.15%   88.14%   -0.01%     
===========================================
  Files          399      399              
  Lines        26372    26410      +38     
  Branches      4114     4122       +8     
===========================================
+ Hits         23247    23278      +31     
- Misses        2241     2245       +4     
- Partials       884      887       +3     

Flag	Coverage Δ
unittests	88.14% <84.09%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
...edit/models/diffusion_schedulers/ddim_scheduler.py	80.85% <ø> (ø)
...edit/models/diffusion_schedulers/ddpm_scheduler.py	75.51% <ø> (ø)
mmedit/registry.py	100.00% <ø> (ø)
mmedit/models/base_archs/__init__.py	91.89% <83.33%> (-8.11%)	⬇️
mmedit/models/diffusion_schedulers/__init__.py	83.33% <83.33%> (ø)
mmedit/models/editors/__init__.py	100.00% <100.00%> (ø)
mmedit/models/editors/ddpm/__init__.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[plyfager]

Seems we didn't support pipelines. How about wrapping pipelines into Inferencer?

[plyfager]

As for training, I support implementing it in BaseModel.train_step.

[zengyh1900]

Seems we didn't support pipelines. How about wrapping pipelines into Inferencer?

The pipelines in diffusers are actually a set of models and process steps, which may be similar to inferences in mmedit.
Since the pipelines can be a bit more complex, I think we can leave the wrappers for pipelines in the future PR.

[zengyh1900]

As for training, I support implementing it in BaseModel.train_step.

Yes. We will support training in BaseModel.train_step.

[LeoXing1996]

LGTM