






"""simple docstring"""




from typing import List, Optional, Tuple, Union

import torch

from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput





class __lowerCamelCase			(				__lowerCAmelCase):





	"""simple docstring"""






	def __init__(							self	,   UpperCAmelCase	,   UpperCAmelCase   ):
							"""simple docstring"""







							super().__init__()
							self.register_modules(unet=__SCREAMING_SNAKE_CASE	,   scheduler=__SCREAMING_SNAKE_CASE   )

	@torch.no_grad()
	def __call__(							self	,   UpperCAmelCase = 1	,   UpperCAmelCase = None	,   UpperCAmelCase = 1000	,   UpperCAmelCase = "pil"	,   UpperCAmelCase = True	,   ):
							"""simple docstring"""







							if isinstance(self.unet.config.sample_size	,   __SCREAMING_SNAKE_CASE   ):
													_UpperCAmelCase     =      (
													    batch_size,
													    self.unet.config.in_channels,
													    self.unet.config.sample_size,
													    self.unet.config.sample_size,
													)
							else:
													_UpperCAmelCase     =      (batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)

							if self.device.type == "mps":
													# randn does not work reproducibly on mps
													_UpperCAmelCase     =      randn_tensor(__SCREAMING_SNAKE_CASE	,   generator=__SCREAMING_SNAKE_CASE   )
													_UpperCAmelCase     =      image.to(self.device   )
							else:
													_UpperCAmelCase     =      randn_tensor(__SCREAMING_SNAKE_CASE	,   generator=__SCREAMING_SNAKE_CASE	,   device=self.device   )

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

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

													# 2. compute previous image: x_t -> x_t-1
													_UpperCAmelCase     =      self.scheduler.step(__SCREAMING_SNAKE_CASE	,   __SCREAMING_SNAKE_CASE	,   __SCREAMING_SNAKE_CASE	,   generator=__SCREAMING_SNAKE_CASE   ).prev_sample

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

							if not return_dict:
													return (image,)

							return ImagePipelineOutput(images=__SCREAMING_SNAKE_CASE   )








import unittest

import torch

from diffusers import DDIMScheduler, TextToVideoZeroPipeline
from diffusers.utils import load_pt, require_torch_gpu, slow

from ..test_pipelines_common import assert_mean_pixel_difference





@slow
@require_torch_gpu
class __lowerCamelCase			(				unittest.TestCase):





	"""simple docstring"""






	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      'runwayml/stable-diffusion-v1-5'
							_UpperCAmelCase     =      TextToVideoZeroPipeline.from_pretrained(UpperCAmelCase	,   torch_dtype=torch.floataa   ).to('cuda'   )
							_UpperCAmelCase     =      DDIMScheduler.from_config(pipe.scheduler.config   )
							_UpperCAmelCase     =      torch.Generator(device='cuda'   ).manual_seed(0   )

							_UpperCAmelCase     =      'A bear is playing a guitar on Times Square'
							_UpperCAmelCase     =      pipe(prompt=UpperCAmelCase	,   generator=UpperCAmelCase   ).images

							_UpperCAmelCase     =      load_pt(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text-to-video/A bear is playing a guitar on Times Square.pt'   )

							assert_mean_pixel_difference(UpperCAmelCase	,   UpperCAmelCase   )
