'''simple docstring'''




import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






class 		__A       (unittest.TestCase   ):






       @property
       def 			_snake_case						(      self      ):
              torch.manual_seed(0      )
              __UpperCAmelCase  :				Dict       				=		UNetaDModel(
                  block_out_channels=(32, 64)				,     layers_per_block=2				,     sample_size=32				,     in_channels=3				,     out_channels=3				,     down_block_types=("DownBlock2D", "AttnDownBlock2D")				,     up_block_types=("AttnUpBlock2D", "UpBlock2D")				,     )
              return model






       @property
       def 			_snake_case						(      self      ):
              torch.manual_seed(0      )
              __UpperCAmelCase  :				int       				=		VQModel(
                  block_out_channels=[32, 64]				,     in_channels=3				,     out_channels=3				,     down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]				,     up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]				,     latent_channels=3				,     )
              return model






       @property
       def 			_snake_case						(      self      ):
              torch.manual_seed(0      )
              __UpperCAmelCase  :				Any       				=		CLIPTextConfig(
                  bos_token_id=0				,     eos_token_id=2				,     hidden_size=32				,     intermediate_size=37				,     layer_norm_eps=1E-05				,     num_attention_heads=4				,     num_hidden_layers=5				,     pad_token_id=1				,     vocab_size=10_00				,     )
              return CLIPTextModel(UpperCamelCase_      )







       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Union[str, Any]       				=		self.dummy_uncond_unet
              __UpperCAmelCase  :				str       				=		DDIMScheduler()
              __UpperCAmelCase  :				Optional[Any]       				=		self.dummy_vq_model

              __UpperCAmelCase  :				Union[str, Any]       				=		LDMPipeline(unet=UpperCamelCase_				,     vqvae=UpperCamelCase_				,     scheduler=UpperCamelCase_      )
              ldm.to(UpperCamelCase_      )
              ldm.set_progress_bar_config(disable=UpperCamelCase_      )

              __UpperCAmelCase  :				Optional[Any]       				=		torch.manual_seed(0      )
              __UpperCAmelCase  :				Dict       				=		ldm(generator=UpperCamelCase_				,     num_inference_steps=2				,     output_type="numpy"      ).images

              __UpperCAmelCase  :				List[str]       				=		torch.manual_seed(0      )
              __UpperCAmelCase  :				Any       				=		ldm(generator=UpperCamelCase_				,     num_inference_steps=2				,     output_type="numpy"				,     return_dict=UpperCamelCase_      )[0]

              __UpperCAmelCase  :				int       				=		image[0, -3:, -3:, -1]
              __UpperCAmelCase  :				Optional[Any]       				=		image_from_tuple[0, -3:, -3:, -1]

              assert image.shape == (1, 64, 64, 3)
              __UpperCAmelCase  :				Optional[int]       				=		np.array([0.8_5_1_2, 0.8_1_8, 0.6_4_1_1, 0.6_8_0_8, 0.4_4_6_5, 0.5_6_1_8, 0.4_6, 0.6_2_3_1, 0.5_1_7_2]      )
              __UpperCAmelCase  :				str       				=		1E-2 if torch_device != "mps" else 3E-2

              assert np.abs(image_slice.flatten() - expected_slice      ).max() < tolerance
              assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < tolerance






@slow
@require_torch
class 		__A       (unittest.TestCase   ):






       def 			_snake_case						(      self      ):
              __UpperCAmelCase  :				Dict       				=		LDMPipeline.from_pretrained("CompVis/ldm-celebahq-256"      )
              ldm.to(UpperCamelCase_      )
              ldm.set_progress_bar_config(disable=UpperCamelCase_      )

              __UpperCAmelCase  :				List[Any]       				=		torch.manual_seed(0      )
              __UpperCAmelCase  :				Dict       				=		ldm(generator=UpperCamelCase_				,     num_inference_steps=5				,     output_type="numpy"      ).images

              __UpperCAmelCase  :				Optional[int]       				=		image[0, -3:, -3:, -1]

              assert image.shape == (1, 2_56, 2_56, 3)
              __UpperCAmelCase  :				List[str]       				=		np.array([0.4_3_9_9, 0.4_4_9_7_5, 0.4_6_8_2_5, 0.4_7_4, 0.4_3_5_9, 0.4_5_8_1, 0.4_5_0_9_5, 0.4_3_4_1, 0.4_4_4_7]      )
              __UpperCAmelCase  :				List[Any]       				=		1E-2 if torch_device != "mps" else 3E-2

              assert np.abs(image_slice.flatten() - expected_slice      ).max() < tolerance





'''simple docstring'''




class 		__A       :






       def __init__(      self				,     UpperCamelCase_      ):
              __UpperCAmelCase  :				Any       				=		set_counts
              __UpperCAmelCase  :				int       				=		max(UpperCamelCase_      )
              __UpperCAmelCase  :				List[str]       				=		len(UpperCamelCase_      )
              __UpperCAmelCase  :				Any       				=		[1] * num_sets
              __UpperCAmelCase  :				Any       				=		list(range(UpperCamelCase_      )      )






       def 			_snake_case						(      self				,     UpperCamelCase_				,     UpperCamelCase_      ):
              __UpperCAmelCase  :				Optional[int]       				=		self.get_parent(UpperCamelCase_      )
              __UpperCAmelCase  :				List[Any]       				=		self.get_parent(UpperCamelCase_      )

              if src_parent == dst_parent:
                     return False

              if self.ranks[dst_parent] >= self.ranks[src_parent]:
                     self.set_counts[dst_parent] += self.set_counts[src_parent]
                     __UpperCAmelCase  :				Optional[Any]       				=		0
                     __UpperCAmelCase  :				List[Any]       				=		dst_parent
                     if self.ranks[dst_parent] == self.ranks[src_parent]:
                            self.ranks[dst_parent] += 1
                     __UpperCAmelCase  :				Union[str, Any]       				=		self.set_counts[dst_parent]
              else:
                     self.set_counts[src_parent] += self.set_counts[dst_parent]
                     __UpperCAmelCase  :				Union[str, Any]       				=		0
                     __UpperCAmelCase  :				Dict       				=		src_parent
                     __UpperCAmelCase  :				Dict       				=		self.set_counts[src_parent]

              __UpperCAmelCase  :				Dict       				=		max(self.max_set				,     UpperCamelCase_      )
              return True







       def 			_snake_case						(      self				,     UpperCamelCase_      ):
              if self.parents[disj_set] == disj_set:
                     return disj_set
              __UpperCAmelCase  :				str       				=		self.get_parent(self.parents[disj_set]      )
              return self.parents[disj_set]






