




'''simple docstring'''


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class    __lowercase						(  __magic_name__       ,							unittest.TestCase ):
   _a							=      UnCLIPImageVariationPipeline
   _a							=      IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
   _a							=      IMAGE_VARIATION_BATCH_PARAMS

   _a							=      [
       """generator""",
       """return_dict""",
       """decoder_num_inference_steps""",
       """super_res_num_inference_steps""",
   ]
   _a							=      False







   @property
   def 		UpperCamelCase__			(	self      )				->    str:
     return 32







   @property
   def 		UpperCamelCase__			(	self      )				->    Optional[int]:
     return 32







   @property
   def 		UpperCamelCase__			(	self      )				->    List[str]:
     return self.time_input_dim







   @property
   def 		UpperCamelCase__			(	self      )				->    Tuple:
     return self.time_input_dim * 4







   @property
   def 		UpperCamelCase__			(	self      )				->    Any:
     return 100







   @property
   def 		UpperCamelCase__			(	self      )				->    Union[str, Any]:
     __a			=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )
     return tokenizer







   @property
   def 		UpperCamelCase__			(	self      )				->    Any:
     torch.manual_seed(0      )
     __a			=					CLIPTextConfig(
         bos_token_id=0      ,							eos_token_id=2      ,							hidden_size=self.text_embedder_hidden_size      ,							projection_dim=self.text_embedder_hidden_size      ,							intermediate_size=37      ,							layer_norm_eps=1e-05      ,							num_attention_heads=4      ,							num_hidden_layers=5      ,							pad_token_id=1      ,							vocab_size=1000      ,							)
     return CLIPTextModelWithProjection(UpperCamelCase      )







   @property
   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     torch.manual_seed(0      )
     __a			=					CLIPVisionConfig(
         hidden_size=self.text_embedder_hidden_size      ,							projection_dim=self.text_embedder_hidden_size      ,							num_hidden_layers=5      ,							num_attention_heads=4      ,							image_size=32      ,							intermediate_size=37      ,							patch_size=1      ,							)
     return CLIPVisionModelWithProjection(UpperCamelCase      )







   @property
   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     torch.manual_seed(0      )

     __a			=					{
         'clip_embeddings_dim': self.text_embedder_hidden_size,
         'time_embed_dim': self.time_embed_dim,
         'cross_attention_dim': self.cross_attention_dim,
     }

     __a			=					UnCLIPTextProjModel(**UpperCamelCase      )
     return model







   @property
   def 		UpperCamelCase__			(	self      )				->    Tuple:
     torch.manual_seed(0      )

     __a			=					{
         'sample_size': 32,
         # RGB in channels
         'in_channels': 3,
         # Out channels is double in channels because predicts mean and variance
         'out_channels': 6,
         'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
         'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
         'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
         'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
         'layers_per_block': 1,
         'cross_attention_dim': self.cross_attention_dim,
         'attention_head_dim': 4,
         'resnet_time_scale_shift': 'scale_shift',
         'class_embed_type': 'identity',
     }

     __a			=					UNetaDConditionModel(**UpperCamelCase      )
     return model







   @property
   def 		UpperCamelCase__			(	self      )				->    List[Any]:
     return {
         "sample_size": 64,
         "layers_per_block": 1,
         "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
         "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
         "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
         "in_channels": 6,
         "out_channels": 3,
     }







   @property
   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     torch.manual_seed(0      )

     __a			=					UNetaDModel(**self.dummy_super_res_kwargs      )
     return model







   @property
   def 		UpperCamelCase__			(	self      )				->    Dict:
     # seeded differently to get different unet than `self.dummy_super_res_first`
     torch.manual_seed(1      )

     __a			=					UNetaDModel(**self.dummy_super_res_kwargs      )
     return model







   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     __a			=					self.dummy_decoder
     __a			=					self.dummy_text_proj
     __a			=					self.dummy_text_encoder
     __a			=					self.dummy_tokenizer
     __a			=					self.dummy_super_res_first
     __a			=					self.dummy_super_res_last

     __a			=					UnCLIPScheduler(
         variance_type='learned_range'      ,							prediction_type='epsilon'      ,							num_train_timesteps=1000      ,							)

     __a			=					UnCLIPScheduler(
         variance_type='fixed_small_log'      ,							prediction_type='epsilon'      ,							num_train_timesteps=1000      ,							)

     __a			=					CLIPImageProcessor(crop_size=32      ,							size=32      )

     __a			=					self.dummy_image_encoder

     return {
         "decoder": decoder,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "text_proj": text_proj,
         "feature_extractor": feature_extractor,
         "image_encoder": image_encoder,
         "super_res_first": super_res_first,
         "super_res_last": super_res_last,
         "decoder_scheduler": decoder_scheduler,
         "super_res_scheduler": super_res_scheduler,
     }







   def 		UpperCamelCase__			(	self      ,							UpperCamelCase      ,							UpperCamelCase=0      ,							UpperCamelCase=True      )				->    Dict:
     __a			=					floats_tensor((1, 3, 32, 32)      ,							rng=random.Random(UpperCamelCase      )      ).to(UpperCamelCase      )
     if str(UpperCamelCase      ).startswith('mps'      ):
       __a			=					torch.manual_seed(UpperCamelCase      )
     else:
       __a			=					torch.Generator(device=UpperCamelCase      ).manual_seed(UpperCamelCase      )

     if pil_image:
       __a			=					input_image * 0.5 + 0.5
       __a			=					input_image.clamp(0      ,							1      )
       __a			=					input_image.cpu().permute(0      ,							2      ,							3      ,							1      ).float().numpy()
       __a			=					DiffusionPipeline.numpy_to_pil(UpperCamelCase      )[0]

     return {
         "image": input_image,
         "generator": generator,
         "decoder_num_inference_steps": 2,
         "super_res_num_inference_steps": 2,
         "output_type": "np",
     }







   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     __a			=					'cpu'

     __a			=					self.get_dummy_components()

     __a			=					self.pipeline_class(**UpperCamelCase      )
     __a			=					pipe.to(UpperCamelCase      )

     pipe.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )

     __a			=					pipe(**UpperCamelCase      )
     __a			=					output.images

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )

     __a			=					pipe(
         **UpperCamelCase      ,							return_dict=UpperCamelCase      ,							)[0]

     __a			=					image[0, -3:, -3:, -1]
     __a			=					image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     __a			=					np.array(
         [
             0.9_997,
             0.0_002,
             0.9_997,
             0.9_997,
             0.9_969,
             0.0_023,
             0.9_997,
             0.9_969,
             0.9_970,
         ]      )

     assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2







   def 		UpperCamelCase__			(	self      )				->    Tuple:
     __a			=					'cpu'

     __a			=					self.get_dummy_components()

     __a			=					self.pipeline_class(**UpperCamelCase      )
     __a			=					pipe.to(UpperCamelCase      )

     pipe.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )

     __a			=					pipe(**UpperCamelCase      )
     __a			=					output.images

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )

     __a			=					pipe(
         **UpperCamelCase      ,							return_dict=UpperCamelCase      ,							)[0]

     __a			=					image[0, -3:, -3:, -1]
     __a			=					image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     __a			=					np.array([0.9_997, 0.0_003, 0.9_997, 0.9_997, 0.9_970, 0.0_024, 0.9_997, 0.9_971, 0.9_971]      )

     assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2







   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     __a			=					'cpu'

     __a			=					self.get_dummy_components()

     __a			=					self.pipeline_class(**UpperCamelCase      )
     __a			=					pipe.to(UpperCamelCase      )

     pipe.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )
     __a			=					[
         pipeline_inputs['image'],
         pipeline_inputs['image'],
     ]

     __a			=					pipe(**UpperCamelCase      )
     __a			=					output.images

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )
     __a			=					[
         tuple_pipeline_inputs['image'],
         tuple_pipeline_inputs['image'],
     ]

     __a			=					pipe(
         **UpperCamelCase      ,							return_dict=UpperCamelCase      ,							)[0]

     __a			=					image[0, -3:, -3:, -1]
     __a			=					image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (2, 64, 64, 3)

     __a			=					np.array(
         [
             0.9_997,
             0.9_989,
             0.0_008,
             0.0_021,
             0.9_960,
             0.0_018,
             0.0_014,
             0.0_002,
             0.9_933,
         ]      )

     assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() < 1e-2







   def 		UpperCamelCase__			(	self      )				->    Union[str, Any]:
     __a			=					torch.device('cpu'      )





     class    __lowercase						:
        _a							=      1

     __a			=					self.get_dummy_components()

     __a			=					self.pipeline_class(**UpperCamelCase      )
     __a			=					pipe.to(UpperCamelCase      )

     pipe.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					torch.Generator(device=UpperCamelCase      ).manual_seed(0      )
     __a			=					pipe.decoder.dtype
     __a			=					1

     __a			=					(
         batch_size,
         pipe.decoder.config.in_channels,
         pipe.decoder.config.sample_size,
         pipe.decoder.config.sample_size,
     )
     __a			=					pipe.prepare_latents(
         UpperCamelCase      ,							dtype=UpperCamelCase      ,							device=UpperCamelCase      ,							generator=UpperCamelCase      ,							latents=UpperCamelCase      ,							scheduler=DummyScheduler()      )

     __a			=					(
         batch_size,
         pipe.super_res_first.config.in_channels // 2,
         pipe.super_res_first.config.sample_size,
         pipe.super_res_first.config.sample_size,
     )
     __a			=					pipe.prepare_latents(
         UpperCamelCase      ,							dtype=UpperCamelCase      ,							device=UpperCamelCase      ,							generator=UpperCamelCase      ,							latents=UpperCamelCase      ,							scheduler=DummyScheduler()      )

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )

     __a			=					pipe(
         **UpperCamelCase      ,							decoder_latents=UpperCamelCase      ,							super_res_latents=UpperCamelCase      ).images

     __a			=					self.get_dummy_inputs(UpperCamelCase      ,							pil_image=UpperCamelCase      )
     # Don't pass image, instead pass embedding
     __a			=					pipeline_inputs.pop('image'      )
     __a			=					pipe.image_encoder(UpperCamelCase      ).image_embeds

     __a			=					pipe(
         **UpperCamelCase      ,							decoder_latents=UpperCamelCase      ,							super_res_latents=UpperCamelCase      ,							image_embeddings=UpperCamelCase      ,							).images

     # make sure passing text embeddings manually is identical
     assert np.abs(img_out_a - img_out_a      ).max() < 1e-4







   @skip_mps
   def 		UpperCamelCase__			(	self      )				->    Dict:
     __a			=					torch_device == 'cpu'

     # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
     __a			=					1e-2

     self._test_attention_slicing_forward_pass(
         test_max_difference=UpperCamelCase      ,							expected_max_diff=UpperCamelCase      )







   @skip_mps
   def 		UpperCamelCase__			(	self      )				->    Tuple:
     __a			=					torch_device == 'cpu'
     __a			=					True
     __a			=					[
         'decoder_num_inference_steps',
         'super_res_num_inference_steps',
     ]

     self._test_inference_batch_single_identical(
         test_max_difference=UpperCamelCase      ,							relax_max_difference=UpperCamelCase      ,							additional_params_copy_to_batched_inputs=UpperCamelCase      ,							)







   def 		UpperCamelCase__			(	self      )				->    Any:
     __a			=					[
         'decoder_num_inference_steps',
         'super_res_num_inference_steps',
     ]

     if torch_device == "mps":
       # TODO: MPS errors with larger batch sizes
       __a			=					[2, 3]
       self._test_inference_batch_consistent(
           batch_sizes=UpperCamelCase      ,							additional_params_copy_to_batched_inputs=UpperCamelCase      ,							)
     else:
       self._test_inference_batch_consistent(
           additional_params_copy_to_batched_inputs=UpperCamelCase      )







   @skip_mps
   def 		UpperCamelCase__			(	self      )				->    List[Any]:
     return super().test_dict_tuple_outputs_equivalent()







   @skip_mps
   def 		UpperCamelCase__			(	self      )				->    List[str]:
     return super().test_save_load_local()







   @skip_mps
   def 		UpperCamelCase__			(	self      )				->    Dict:
     return super().test_save_load_optional_components()





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







   def 		UpperCamelCase__			(	self      )				->    Tuple:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()







   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     __a			=					load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png'      )
     __a			=					load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/unclip/karlo_v1_alpha_cat_variation_fp16.npy'      )

     __a			=					UnCLIPImageVariationPipeline.from_pretrained(
         'kakaobrain/karlo-v1-alpha-image-variations'      ,							torch_dtype=torch.floataa      )
     __a			=					pipeline.to(UpperCamelCase      )
     pipeline.set_progress_bar_config(disable=UpperCamelCase      )

     __a			=					torch.Generator(device='cpu'      ).manual_seed(0      )
     __a			=					pipeline(
         UpperCamelCase      ,							generator=UpperCamelCase      ,							output_type='np'      ,							)

     __a			=					output.images[0]

     assert image.shape == (256, 256, 3)

     assert_mean_pixel_difference(UpperCamelCase      ,							UpperCamelCase      ,							15      )









'''simple docstring'''


import torch

from diffusers import StableDiffusionPipeline


UpperCAmelCase_ 		=      "path-to-your-trained-model"
UpperCAmelCase_ 		=      StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.floataa).to("cuda")

UpperCAmelCase_ 		=      "A photo of sks dog in a bucket"
UpperCAmelCase_ 		=      pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save("dog-bucket.png")



