[WIP] Add UFOGen Pipeline and Scheduler

[dg845]

What does this PR do?

This PR adds a pipeline and scheduler for the UFOGen model. The UFOGen model is based on the denoising diffusion GAN (DDGAN) model with modifications to enable one-step sampling.

The GAN discriminator used for adversarial training of a UFOGen model is not included in this PR.

Partially addresses #5905. See also #5979.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@patrickvonplaten
@patil-suraj
@JunbongJang

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[dg845]

The original UFOGen paper doesn't explicitly give an algorithm for multistep sampling (and there is no official implementation currently available AFAIK); the current multistep sampling implementation is my current best guess at what the right algorithm should be.

Right now, for sampling steps which are not the final sampling step, we get the pred_prev_sample $x_{t - 1}$ by sampling from $q(x_{t - 1} \mid x_0 = G_\theta(x_t, t))$, where $q(x_{t - 1} \mid x_0) = \mathcal{N}(\sqrt{\bar{\alpha}_{t - 1}}\boldsymbol{x}_0, (1 - \bar{\alpha}_{t - 1})\boldsymbol{I})$ is the true forward process, following the parameterization of the U-Net generator in Section 4 of the paper. (This is very similar to the distribution we sample from when we add_noise using the scheduler, but at timestep $t - 1$ instead of timestep $t$.)

For the final sampling step, or for one-step sampling, we simply return the pred_original_sample $\hat{\boldsymbol{x}}_0 = G_\theta(\boldsymbol{x}_t, t)$ without any sampling.

[dg845]

Since UFOGen models typically use the same architecture and initial weights as a pretrained diffusion model (and in particular Stable Diffusion v1.5, see sections 4.2 and 5.1 of the paper), I think it probably makes sense to not have a separate UFOGenPipeline and instead make sure UFOGenScheduler is compatible with existing diffusion model pipelines like StableDiffusionPipeline, similar to how LCM models are supported. Thoughts? @patrickvonplaten @patil-suraj

(I have preliminarily added a test in tests/pipelines/stable_diffusion/test_stable_diffusion.py to test compatibility between StableDiffusionPipeline and UFOGenScheduler.)

[dg845]

Assuming the multistep sampling strategy described in #6133 (comment) is correct, the UFOGenScheduler ends up being very similar to LCMScheduler, as both resolve the predicted_original_sample $\boldsymbol{x}_0$ and then sample from the true forward process

$$q(\boldsymbol{x}_{t_{prev}} \mid \boldsymbol{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_{t_{prev}}}\boldsymbol{x}_0, (1 - \bar{\alpha}_{t_{prev}})\boldsymbol{I})$$

at the next scheduled timestep $t_{prev}$ when we're not at the last step of the timestep schedule:

diffusers/src/diffusers/schedulers/scheduling_lcm.py

Lines 554 to 558 in 1d686ba

	if self.step_index != self.num_inference_steps - 1:
	noise = randn_tensor(
	model_output.shape, generator=generator, device=model_output.device, dtype=denoised.dtype
	)
	prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise

So it could be the case that a separate UFOGenScheduler is not strictly necessary. However, there are several differences between LCMScheduler and UFOGenScheduler:

• LCM models are parameterized in a special way to enforce the consistency model boundary conditions:

diffusers/src/diffusers/schedulers/scheduling_lcm.py

Lines 548 to 549 in 1d686ba

	# 6. Denoise model output using boundary conditions
	denoised = c_out * predicted_original_sample + c_skip * sample

whereas UFOGen models do not need to be parameterized in this way (and so UFOGenScheduler does not implement this logic).

• The LCMScheduler.set_timesteps method will by default only select a subset of the timesteps on the original training/distillation timestep schedule, while UFOGenScheduler.set_timesteps currently doesn't implement this restriction (and is very similar to DDPMScheduler.set_timesteps).

[sayakpaul]

@dg845 is this ready for a review?

[dg845]

Yes (although as far as I know there is no official implementation to compare the PR to, and I'm not sure if there are any publicly available UFOGen checkpoints to use with UFOGenScheduler currently).

[importnumpy]

I have one question. Do you have any plans for UFOGen training code yet?
I'm curious about how to actually test the current PR.

[dg845]

Hi @importnumpy, @patil-suraj and I are working on a UFOGen training script, but it is not ready yet.

[github-actions]

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.

[patrickvonplaten]

Should we maybe instead add this pipeline to the community folder

[dg845]

Sounds good to me. I think the main blocker for the PR is that AFAIK there are no strong UFOGen checkpoints to test the pipeline and scheduler with.

[github-actions]

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.

[sayakpaul]

I don’t think this is stale.

[yiyixuxu]

@dg845
are we still working on this? I saw that UFOGen scheduler PR is already merged.

[dg845]

Hi @yiyixuxu, sorry for the late reply. I am not actively working on this but I think it makes sense to keep this PR around in case an open source UFOGen code repository/checkpoint comes out. In that case, we could test this PR against the reference code/checkpoint with StableDiffusionPipeline + UFOGenScheduler (if it reproduces the Stable Diffusion v1.5 setup used in the UFOGen paper) and potentially merge it as an "official" UFOGen scheduler.

[weleen]

Hi @yiyixuxu, sorry for the late reply. I am not actively working on this but I think it makes sense to keep this PR around in case an open source UFOGen code repository/checkpoint comes out. In that case, we could test this PR against the reference code/checkpoint with StableDiffusionPipeline + UFOGenScheduler (if it reproduces the Stable Diffusion v1.5 setup used in the UFOGen paper) and potentially merge it as an "official" UFOGen scheduler.

Hi @dg845, Would you please share the link to the open-source UFOGen repository for further testing?

[dg845]

Hi @weleen, to the best of my knowledge there are no (official or unofficial) open source UFOGen implementations currently available.

[weleen]

@dg845 Thanks!

[github-actions]

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.