hf-internal-testing
/

diffusers-dummy-pipeline

Model card Files Files and versions Community

diffusers-dummy-pipeline / pipeline.py

patrickvonplaten

Update pipeline.py

278ec4f about 1 year ago

raw history blame contribute delete

No virus

4.29 kB

	# Copyright 2022 The HuggingFace Team. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and

	# limitations under the License.


	from typing import Optional, Tuple, Union

	import torch

	from diffusers import DiffusionPipeline, ImagePipelineOutput


	class CustomPipeline(DiffusionPipeline):
	r"""
	This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
	library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

	Parameters:
	unet ([`UNet2DModel`]): U-Net architecture to denoise the encoded image.
	scheduler ([`SchedulerMixin`]):
	A scheduler to be used in combination with `unet` to denoise the encoded image. Can be one of
	[`DDPMScheduler`], or [`DDIMScheduler`].
	"""

	def __init__(self, unet, scheduler):
	super().__init__()
	self.register_modules(unet=unet, scheduler=scheduler)

	@torch.no_grad()
	def __call__(
	self,
	batch_size: int = 1,
	generator: Optional[torch.Generator] = None,
	num_inference_steps: int = 50,
	output_type: Optional[str] = "pil",
	return_dict: bool = True,
	**kwargs,
	) -> Union[ImagePipelineOutput, Tuple]:
	r"""
	Args:
	batch_size (`int`, optional, defaults to 1):
	The number of images to generate.
	generator (`torch.Generator`, optional):
	A [torch generator](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation
	deterministic.
	eta (`float`, optional, defaults to 0.0):
	The eta parameter which controls the scale of the variance (0 is DDIM and 1 is one type of DDPM).
	num_inference_steps (`int`, optional, defaults to 50):
	The number of denoising steps. More denoising steps usually lead to a higher quality image at the
	expense of slower inference.
	output_type (`str`, optional, defaults to `"pil"`):
	The output format of the generate image. Choose between
	[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
	return_dict (`bool`, optional, defaults to `True`):
	Whether or not to return a [`~pipeline_utils.ImagePipelineOutput`] instead of a plain tuple.

	Returns:
	[`~pipeline_utils.ImagePipelineOutput`] or `tuple`: [`~pipelines.utils.ImagePipelineOutput`] if
	`return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the
	generated images.
	"""

	# Sample gaussian noise to begin loop
	image = torch.randn(
	(batch_size, self.unet.config.in_channels, self.unet.config.sample_size, self.unet.config.sample_size),
	generator=generator,
	)
	image = image.to(self.device)

	# set step values
	self.scheduler.set_timesteps(num_inference_steps)

	for t in self.progress_bar(self.scheduler.timesteps):
	# 1. predict noise model_output
	model_output = self.unet(image, t).sample

	# 2. predict previous mean of image x_t-1 and add variance depending on eta
	# eta corresponds to η in paper and should be between [0, 1]
	# do x_t -> x_t-1
	image = self.scheduler.step(model_output, t, image).prev_sample

	image = (image / 2 + 0.5).clamp(0, 1)
	image = image.cpu().permute(0, 2, 3, 1).numpy()
	if output_type == "pil":
	image = self.numpy_to_pil(image)

	if not return_dict:
	return (image,)

	return ImagePipelineOutput(images=image), "This is a test"