diffusers icon indicating copy to clipboard operation
diffusers copied to clipboard

New modalities

Open richardrl opened this issue 3 years ago • 6 comments

One of the stated design decisions from the readme was to support arbitrary modalities, not just images.

I'm in the process of trying to adapt the code for 1D vectors (not H x W x C images).

this line:

noisy_images = noise_scheduler.training_step(clean_inputs, noise_samples, timesteps)

adapted from train_unconditional.py, takes in a (16, 198) clean_inputs tensor and returns a (16, 1, 16, 198) noisy_images tensor. So, 1D tensors are not working out of the box. I am just curious if I am taking the right approach to get 1D tensors to work or there's no avoiding custom coding a new DDPMScheduler etc.

EDIT: I got the DDPM training_step function to work with this change:

        if len(original_samples.shape) == 2:
            timesteps = timesteps.unsqueeze(-1)
        else:
            timesteps = timesteps.reshape(batch_size, 1, 1, 1)

richardrl avatar Jun 24 '22 04:06 richardrl

Thanks for reporting @richardrl ! Indeed the plan is to support multiple modalities, but we haven't yet tested the schedulers with 1D data.

cc @patrickvonplaten @anton-l

patil-suraj avatar Jun 24 '22 09:06 patil-suraj

@anton-l - we need to make sure that training_step is both framework agnostic and shape agnostic

patrickvonplaten avatar Jun 24 '22 12:06 patrickvonplaten

Taking a look at the function as is:

    def training_step(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor):
        if timesteps.dim() != 1:
            raise ValueError("`timesteps` must be a 1D tensor")

        device = original_samples.device
        batch_size = original_samples.shape[0]
        timesteps = timesteps.reshape(batch_size, 1, 1, 1)

        sqrt_alpha_prod = self.alphas_cumprod[timesteps] ** 0.5
        sqrt_one_minus_alpha_prod = (1 - self.alphas_cumprod[timesteps]) ** 0.5
        noisy_samples = sqrt_alpha_prod.to(device) * original_samples + sqrt_one_minus_alpha_prod.to(device) * noise
        return noisy_samples

Note that the input can be both torch and numpy tensors -> this should be changed.

Also there shouldn't be any .to(device) statements, nor framework and modality spefific .reshape(...) operation.

I'd be in favor of implementing framework specific (one for PT one for TF) functions called

def extract(....) in SchedulerMixin that have if framework == "pt" statements in them. Also note that we shouldn't assume to know the dimension of the input original_samples

patrickvonplaten avatar Jun 24 '22 12:06 patrickvonplaten

BTW, it's super nice to get all your feedback here @richardrl - thanks a lot!

patrickvonplaten avatar Jun 24 '22 13:06 patrickvonplaten

Now we use match_shape(timestaps, original_samples) for everything, which is framework- and shape-agnostic: https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_ddpm.py#L146

anton-l avatar Jun 27 '22 15:06 anton-l

@anton-l @patrickvonplaten Thanks for your input thus far.

I took the latest commit (as of this moment) and made a minimum reproduction of a 1D MLP model and training.

I had to make an additional modification to pipeline_ddpm.py to support noise samples of the right shape.

bimodal_testt.zip

python3 train_unconditional_gaussian_test.py This runs a test on a bimodal gaussian distribution centered at +33, -33 with low variance

It seems to not capture the -33 mode after an epoch or two. Am running the training overnight to see what happens.

Welcome you guys to try running this to see if there's anything I did wrong

richardrl avatar Jun 28 '22 08:06 richardrl

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

github-actions[bot] avatar Sep 16 '22 15:09 github-actions[bot]

@anton-l as you've re-opened the issue -> are you planning on doing something with it?

patrickvonplaten avatar Oct 27 '22 08:10 patrickvonplaten

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

github-actions[bot] avatar Nov 20 '22 15:11 github-actions[bot]

I think we have support for all shapes now, agreed with the stalebot :)

anton-l avatar Nov 28 '22 15:11 anton-l